|
1 | 1 | import ast
|
| 2 | +import importlib |
2 | 3 | import inspect
|
| 4 | +import os |
3 | 5 | import textwrap
|
4 | 6 | from collections import defaultdict
|
5 | 7 | from typing import Literal, cast
|
@@ -123,107 +125,55 @@ def visit_FunctionDef(self, node):
|
123 | 125 | return ast.copy_location(new_node, node)
|
124 | 126 |
|
125 | 127 |
|
126 |
| -from weaviate.collections.aggregations.hybrid import executor as agg_hybrid |
127 |
| -from weaviate.collections.aggregations.near_image import executor as agg_near_image |
128 |
| -from weaviate.collections.aggregations.near_object import executor as agg_near_object |
129 |
| -from weaviate.collections.aggregations.near_text import executor as agg_near_text |
130 |
| -from weaviate.collections.aggregations.near_vector import executor as agg_near_vector |
131 |
| -from weaviate.collections.aggregations.over_all import executor as agg_over_all |
132 |
| -from weaviate.collections.backups import executor as backups |
133 |
| -from weaviate.collections.cluster import executor as cluster |
134 |
| -from weaviate.collections.config import executor as config |
135 |
| -from weaviate.collections.data import executor as data |
136 |
| -from weaviate.collections.queries.bm25.generate import executor as generate_bm25 |
137 |
| -from weaviate.collections.queries.bm25.query import executor as query_bm25 |
138 |
| -from weaviate.collections.queries.fetch_object_by_id import executor as fetch_object_by_id |
139 |
| -from weaviate.collections.queries.fetch_objects.generate import executor as generate_fetch_objects |
140 |
| -from weaviate.collections.queries.fetch_objects.query import executor as query_fetch_objects |
141 |
| -from weaviate.collections.queries.fetch_objects_by_ids.generate import ( |
142 |
| - executor as generate_fetch_objects_by_ids, |
143 |
| -) |
144 |
| -from weaviate.collections.queries.fetch_objects_by_ids.query import ( |
145 |
| - executor as query_fetch_objects_by_ids, |
146 |
| -) |
147 |
| -from weaviate.collections.queries.hybrid.generate import executor as generate_hybrid |
148 |
| -from weaviate.collections.queries.hybrid.query import executor as query_hybrid |
149 |
| -from weaviate.collections.queries.near_image.generate import executor as generate_near_image |
150 |
| -from weaviate.collections.queries.near_image.query import executor as query_near_image |
151 |
| -from weaviate.collections.queries.near_media.generate import executor as generate_near_media |
152 |
| -from weaviate.collections.queries.near_media.query import executor as query_near_media |
153 |
| -from weaviate.collections.queries.near_object.generate import executor as generate_near_object |
154 |
| -from weaviate.collections.queries.near_object.query import executor as query_near_object |
155 |
| -from weaviate.collections.queries.near_text.generate import executor as generate_near_text |
156 |
| -from weaviate.collections.queries.near_text.query import executor as query_near_text |
157 |
| -from weaviate.collections.queries.near_vector.generate import executor as generate_near_vector |
158 |
| -from weaviate.collections.queries.near_vector.query import executor as query_near_vector |
159 |
| -from weaviate.debug import executor as debug |
160 |
| -from weaviate.rbac import executor as rbac |
161 |
| -from weaviate.collections.tenants import executor as tenants |
162 |
| -from weaviate.users import executor as users |
163 |
| - |
164 |
| -for module in [ |
165 |
| - agg_hybrid, |
166 |
| - agg_near_image, |
167 |
| - agg_near_object, |
168 |
| - agg_near_text, |
169 |
| - agg_near_vector, |
170 |
| - agg_over_all, |
171 |
| - backups, |
172 |
| - cluster, |
173 |
| - config, |
174 |
| - data, |
175 |
| - debug, |
176 |
| - generate_bm25, |
177 |
| - generate_fetch_objects, |
178 |
| - generate_fetch_objects_by_ids, |
179 |
| - generate_hybrid, |
180 |
| - generate_near_image, |
181 |
| - generate_near_media, |
182 |
| - generate_near_object, |
183 |
| - generate_near_text, |
184 |
| - generate_near_vector, |
185 |
| - fetch_object_by_id, |
186 |
| - query_bm25, |
187 |
| - query_fetch_objects, |
188 |
| - query_fetch_objects_by_ids, |
189 |
| - query_hybrid, |
190 |
| - query_near_image, |
191 |
| - query_near_media, |
192 |
| - query_near_object, |
193 |
| - query_near_text, |
194 |
| - query_near_vector, |
195 |
| - rbac, |
196 |
| - tenants, |
197 |
| - users, |
198 |
| -]: |
199 |
| - source = textwrap.dedent(inspect.getsource(module)) |
200 |
| - |
201 |
| - colours: list[Literal["sync", "async"]] = ["sync", "async"] |
202 |
| - for colour in colours: |
203 |
| - tree = ast.parse(source, mode="exec", type_comments=True) |
204 |
| - |
205 |
| - transformer = ExecutorTransformer(colour) |
206 |
| - stubbed = transformer.visit(tree) |
207 |
| - |
208 |
| - imports = [ |
209 |
| - node for node in stubbed.body if isinstance(node, (ast.Import, ast.ImportFrom)) |
210 |
| - ] + [ |
211 |
| - ast.ImportFrom( |
212 |
| - module="weaviate.connect.v4", |
213 |
| - names=[ast.alias(name=f"Connection{colour.capitalize()}", asname=None)], |
214 |
| - level=0, |
215 |
| - ), |
216 |
| - ast.ImportFrom( |
217 |
| - module=".executor", |
218 |
| - names=[ast.alias(name=name, asname=None) for name in transformer.executor_names], |
219 |
| - level=0, |
220 |
| - ), |
221 |
| - ] |
222 |
| - stubbed.body = imports + [node for node in stubbed.body if isinstance(node, ast.ClassDef)] |
223 |
| - ast.fix_missing_locations(stubbed) |
224 |
| - |
225 |
| - dir = cast(str, module.__package__).replace(".", "/") |
226 |
| - file = f"{dir}/{colour}.pyi" if colour == "sync" else f"{dir}/{colour}_.pyi" |
227 |
| - with open(file, "w") as f: |
228 |
| - print(f"Writing {file}") |
229 |
| - f.write(ast.unparse(stubbed)) |
| 128 | +for subdir, dirs, files in os.walk("./weaviate"): |
| 129 | + for file in files: |
| 130 | + if file != "executor.py": |
| 131 | + continue |
| 132 | + if "connect" in subdir: |
| 133 | + # ignore weaviate/connect/executor.py file |
| 134 | + continue |
| 135 | + if "collections/collections" in subdir: |
| 136 | + # ignore weaviate/collections/collections directory |
| 137 | + continue |
| 138 | + |
| 139 | + mod = os.path.join(subdir, file) |
| 140 | + mod = mod[2:] # remove the leading dot and slash |
| 141 | + mod = mod[:-3] # remove the .py |
| 142 | + mod = mod.replace("/", ".") # convert into pythonic import |
| 143 | + |
| 144 | + module = importlib.import_module(mod) |
| 145 | + source = textwrap.dedent(inspect.getsource(module)) |
| 146 | + |
| 147 | + colours: list[Literal["sync", "async"]] = ["sync", "async"] |
| 148 | + for colour in colours: |
| 149 | + tree = ast.parse(source, mode="exec", type_comments=True) |
| 150 | + |
| 151 | + transformer = ExecutorTransformer(colour) |
| 152 | + stubbed = transformer.visit(tree) |
| 153 | + |
| 154 | + imports = [ |
| 155 | + node for node in stubbed.body if isinstance(node, (ast.Import, ast.ImportFrom)) |
| 156 | + ] + [ |
| 157 | + ast.ImportFrom( |
| 158 | + module="weaviate.connect.v4", |
| 159 | + names=[ast.alias(name=f"Connection{colour.capitalize()}", asname=None)], |
| 160 | + level=0, |
| 161 | + ), |
| 162 | + ast.ImportFrom( |
| 163 | + module=".executor", |
| 164 | + names=[ |
| 165 | + ast.alias(name=name, asname=None) for name in transformer.executor_names |
| 166 | + ], |
| 167 | + level=0, |
| 168 | + ), |
| 169 | + ] |
| 170 | + stubbed.body = imports + [ |
| 171 | + node for node in stubbed.body if isinstance(node, ast.ClassDef) |
| 172 | + ] |
| 173 | + ast.fix_missing_locations(stubbed) |
| 174 | + |
| 175 | + dir = cast(str, module.__package__).replace(".", "/") |
| 176 | + file = f"{dir}/{colour}.pyi" if colour == "sync" else f"{dir}/{colour}_.pyi" |
| 177 | + with open(file, "w") as f: |
| 178 | + print(f"Writing {file}") |
| 179 | + f.write(ast.unparse(stubbed)) |
0 commit comments