|
| 1 | +import ast |
| 2 | +import importlib |
| 3 | +import inspect |
| 4 | +import os |
| 5 | +import textwrap |
| 6 | +from collections import defaultdict |
| 7 | +from typing import Literal, cast |
| 8 | + |
| 9 | + |
| 10 | +class ExecutorTransformer(ast.NodeTransformer): |
| 11 | + def __init__(self, colour: Literal["async", "sync"]): |
| 12 | + self.colour = colour |
| 13 | + self.executor_names = [] |
| 14 | + |
| 15 | + def visit_ClassDef(self, node): |
| 16 | + self.executor_names.append(node.name) |
| 17 | + node.bases = self.__parse_generics(node) |
| 18 | + node.body = self.__parse_body(node) |
| 19 | + node.name = node.name.replace( |
| 20 | + "Executor", "" if self.colour == "sync" else self.colour.capitalize() |
| 21 | + ) |
| 22 | + self.generic_visit(node) |
| 23 | + return node |
| 24 | + |
| 25 | + def __is_overload(self, fn: ast.FunctionDef): |
| 26 | + return any(isinstance(d, ast.Name) and d.id == "overload" for d in fn.decorator_list) |
| 27 | + |
| 28 | + def __parse_body(self, node: ast.ClassDef): |
| 29 | + funcs_by_name = defaultdict(list) |
| 30 | + for stmt in node.body: |
| 31 | + if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)): |
| 32 | + funcs_by_name[stmt.name].append(stmt) |
| 33 | + |
| 34 | + new_body: list[ast.stmt] = [] |
| 35 | + for stmt in node.body: |
| 36 | + if isinstance(stmt, ast.FunctionDef) and stmt.name.startswith("__"): |
| 37 | + continue # Skip all dunder methods |
| 38 | + if isinstance(stmt, ast.FunctionDef): |
| 39 | + overloads = funcs_by_name[stmt.name] |
| 40 | + if any(self.__is_overload(f) for f in overloads): |
| 41 | + if not self.__is_overload(stmt): |
| 42 | + continue # skip the impl |
| 43 | + new_body.append(stmt) |
| 44 | + return new_body |
| 45 | + |
| 46 | + def __parse_generics(self, node: ast.ClassDef): |
| 47 | + new_bases: list[ast.expr] = [] |
| 48 | + for base in node.bases: |
| 49 | + if not isinstance(base, ast.Subscript): |
| 50 | + continue |
| 51 | + if isinstance(base.value, ast.Name) and base.value.id == "Generic": |
| 52 | + # This is a generic class |
| 53 | + # We need to extract the type arguments |
| 54 | + if isinstance(base.slice, ast.Tuple): |
| 55 | + # This is a tuple of types |
| 56 | + # must remove `ConnectionType` if there |
| 57 | + generics = [ |
| 58 | + arg.id |
| 59 | + for arg in base.slice.elts |
| 60 | + if isinstance(arg, ast.Name) |
| 61 | + if arg.id != "ConnectionType" |
| 62 | + ] |
| 63 | + new_bases.append( |
| 64 | + ast.Subscript( |
| 65 | + value=base.value, |
| 66 | + slice=ast.Tuple( |
| 67 | + elts=[ast.Name(id=arg) for arg in generics], ctx=ast.Load() |
| 68 | + ), |
| 69 | + ctx=ast.Load(), |
| 70 | + ) |
| 71 | + ) |
| 72 | + elif isinstance(base.slice, ast.Name): |
| 73 | + # This is a single type |
| 74 | + if base.slice.id == "ConnectionType": |
| 75 | + # We don't want to include ConnectionType |
| 76 | + continue |
| 77 | + new_bases.append(base) |
| 78 | + connection_type = ast.Name(id=self.__which_connection_type(), ctx=ast.Load()) |
| 79 | + if len(new_bases) == 0: |
| 80 | + # no generics, we need to add the ConnectionType |
| 81 | + slice = connection_type |
| 82 | + else: |
| 83 | + elts: list[ast.expr] = [] |
| 84 | + for base in new_bases: |
| 85 | + assert isinstance(base, ast.Subscript) |
| 86 | + slice = base.slice |
| 87 | + assert isinstance(slice, ast.Tuple) |
| 88 | + elts.extend(slice.elts) |
| 89 | + slice = ast.Tuple(elts=[connection_type, *elts], ctx=ast.Load()) |
| 90 | + new_bases.append( |
| 91 | + ast.Subscript( |
| 92 | + value=ast.Name(id=node.name, ctx=ast.Load()), |
| 93 | + slice=slice, |
| 94 | + ctx=ast.Load(), |
| 95 | + ) |
| 96 | + ) |
| 97 | + return new_bases |
| 98 | + |
| 99 | + def __which_connection_type(self): |
| 100 | + return "ConnectionAsync" if self.colour == "async" else "ConnectionSync" |
| 101 | + |
| 102 | + def __extract_inner_return_type(self, node: ast.expr | None) -> ast.expr | None: |
| 103 | + # Looking for executor.Result[T] |
| 104 | + if ( |
| 105 | + isinstance(node, ast.Subscript) |
| 106 | + and isinstance(node.value, ast.Attribute) |
| 107 | + and isinstance(node.value.value, ast.Name) |
| 108 | + and node.value.value.id == "executor" |
| 109 | + and node.value.attr == "Result" |
| 110 | + ): |
| 111 | + # This is executor.Result[...] |
| 112 | + return node.slice # Return T |
| 113 | + return node # fallback, return original if not matching |
| 114 | + |
| 115 | + def visit_FunctionDef(self, node): |
| 116 | + func_def = ast.AsyncFunctionDef if self.colour == "async" else ast.FunctionDef |
| 117 | + new_node = func_def( |
| 118 | + name=node.name, |
| 119 | + args=node.args, |
| 120 | + body=[ast.Expr(value=ast.Constant(value=Ellipsis))], |
| 121 | + decorator_list=node.decorator_list, |
| 122 | + returns=self.__extract_inner_return_type(node.returns), |
| 123 | + type_comment=node.type_comment, |
| 124 | + ) |
| 125 | + return ast.copy_location(new_node, node) |
| 126 | + |
| 127 | + |
| 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