add script to create stubs from yaml

integration/test_collection.py

@@ -604,16 +604,15 @@ def test_search_hybrid(collection_factory: CollectionFactory, fusion_type: Hybri
     )
     collection.data.insert({"Name": "some name"}, uuid=uuid.uuid4())
     collection.data.insert({"Name": "other word"}, uuid=uuid.uuid4())
-    objs = collection.query.hybrid(
+    objs1 = collection.query.hybrid(
         alpha=0, query="name", fusion_type=fusion_type, include_vector=True
     ).objects
-    assert len(objs) == 1
+    assert len(objs1) == 1
 
-    assert objs[0].vector is not None
-    objs = collection.query.hybrid(
-        alpha=1, query="name", fusion_type=fusion_type, vector=objs[0].vector["default"]
+    objs2 = collection.query.hybrid(
+        alpha=1, query="name", fusion_type=fusion_type, vector=objs1[0].vector["default"]
     ).objects
-    assert len(objs) == 2
+    assert len(objs2) == 2
 
 
 @pytest.mark.skip(reason="currently bugged in weaviate")
@@ -1723,11 +1722,11 @@ class DataModel4(TypedDict):
         pass
 
     objects: Union[
-        List[Object[DataModel0, None]],
-        List[Object[DataModel1, None]],
-        List[Object[DataModel2, None]],
-        List[Object[DataModel3, None]],
-        List[Object[DataModel4, None]],
+        List[Object[DataModel0, None, None]],
+        List[Object[DataModel1, None, None]],
+        List[Object[DataModel2, None, None]],
+        List[Object[DataModel3, None, None]],
+        List[Object[DataModel4, None, None]],
     ]
     if which_case == 0:
         objects = collection.query.fetch_objects(return_properties=DataModel0).objects

integration/test_collection_references.py

@@ -367,7 +367,7 @@ class CRefs(TypedDict):
         )
     elif level == "col-query":
         c_objs = (
-            collection_factory_get(C.name, CProps)
+            collection_factory_get(C.name, CProps)  # type: ignore
             .query.bm25(
                 query="find",
                 include_vector=True,
@@ -387,7 +387,7 @@ class CRefs(TypedDict):
         )
     else:
         c_objs = (
-            collection_factory_get(C.name)
+            collection_factory_get(C.name)  # type: ignore
             .query.bm25(
                 query="find",
                 include_vector=True,
@@ -405,6 +405,7 @@ class CRefs(TypedDict):
     assert (
         c_objs[0].properties.get("not_specified") is None
     )  # type is str but instance is None (in type but not in return_properties)
+    assert c_objs[0].references is not None
     assert c_objs[0].references["b"].objects[0].collection == B.name
     assert c_objs[0].references["b"].objects[0].properties["name"] == "B"
     assert c_objs[0].references["b"].objects[0].uuid == uuid_B

integration/test_iterator.py

@@ -69,7 +69,9 @@ def test_iterator_arguments(
     )
 
     iter_ = collection.iterator(
-        include_vector, return_metadata=return_metadata, return_properties=return_properties
+        include_vector=include_vector,
+        return_metadata=return_metadata,
+        return_properties=return_properties,
     )
 
     # Expect everything back

tools/stubs.py

@@ -0,0 +1,75 @@
+import yaml
+from typing import List
+
+
+def parse_yaml(file_path):
+    with open(file_path, "r") as file:
+        return yaml.safe_load(file)
+
+
+class Function:
+    def __init__(self, name: str, args: List[str], return_type: str) -> None:
+        self.name = name
+        self.args = args
+        self.return_type = return_type
+
+    def __str__(self) -> str:
+        return f"def {self.name}({', '.join(self.args)}) -> {self.return_type}: ..."
+
+
+def generate_pyi(data: dict):
+    for function in data["functions"]:
+        functions: List[Function] = []
+        required: List[str] = []
+        for arg in function["required"]:
+            if arg["name"] == "self":
+                required.append("self")
+            else:
+                required.append(f"{arg['name']}: {arg['type']}")
+
+        optional: List[str] = []
+        if function.get("optional") is not None:
+            for arg in function["optional"]:
+                optional.append(f"{arg['name']}: {arg['type']}")
+
+        args = required + ["*"] + optional
+
+        for i_v in function["include_vector"]:
+            for r_p in function["return_properties"]:
+                for r_r in function["return_references"]:
+                    if function.get("group_by") is not None:
+                        for group in function["group_by"]:
+                            argss = args.copy()
+                            argss.append(f"group_by: {group['type']}")
+                            argss.append(f"include_vector: {i_v['type']}")
+                            argss.append(f"return_properties: {r_p['type']}")
+                            argss.append(f"return_references: {r_r['type']}")
+                            if i_v["generic"] == "None,Vectors":
+                                return_ = f"Union[{group['return']}[{r_p['generic']}, {r_r['generic']}, None], {group['return']}[{r_p['generic']}, {r_r['generic']}, Vectors]]"
+                            else:
+                                return_ = f"{group['return']}[{r_p['generic']}, {r_r['generic']}, {i_v['generic']}]"
+                            functions.append(
+                                Function(name=function["name"], args=argss, return_type=return_)
+                            )
+                    else:
+                        argss = args.copy()
+                        argss.append(f"include_vector: {i_v['type']}")
+                        argss.append(f"return_properties: {r_p['type']}")
+                        argss.append(f"return_references: {r_r['type']}")
+                        if i_v["generic"] == "None,Vectors":
+                            return_ = f"Union[{function['return']}[{r_p['generic']}, {r_r['generic']}, None], {function['return']}[{r_p['generic']}, {r_r['generic']}, Vectors]]"
+                        else:
+                            return_ = f"{function['return']}[{r_p['generic']}, {r_r['generic']}, {i_v['generic']}]"
+                        functions.append(
+                            Function(name=function["name"], args=argss, return_type=return_)
+                        )
+
+        with open(function["target_file"], "w") as file:
+            file.write("\n".join(function["header"]))
+            for func in functions:
+                file.write(f"\n    @overload\n    {func}\n")
+
+
+# Example usage
+yaml_data = parse_yaml("./tools/stubs.yaml")
+generate_pyi(yaml_data)