Add type hints to parsers

Author: GlassOfWhiskey

schema_salad/codegen_base.py

@@ -23,6 +23,7 @@ class LazyInitDef(NamedTuple):
 
     name: str
     init: str
+    instance_type: str | None = None
 
 
 class CodeGenBase:

schema_salad/python_codegen.py

@@ -17,12 +17,18 @@
 from .exceptions import SchemaException
 from .schema import shortname
 
-_string_type_def: Final = TypeDef("strtype", "_PrimitiveLoader(str)")
-_int_type_def: Final = TypeDef("inttype", "_PrimitiveLoader(int)")
-_float_type_def: Final = TypeDef("floattype", "_PrimitiveLoader(float)")
-_bool_type_def: Final = TypeDef("booltype", "_PrimitiveLoader(bool)")
-_null_type_def: Final = TypeDef("None_type", "_PrimitiveLoader(type(None))")
-_any_type_def: Final = TypeDef("Any_type", "_AnyLoader()")
+_string_type_def: Final = TypeDef(name="strtype", init="_PrimitiveLoader(str)", instance_type="str")
+_int_type_def: Final = TypeDef(name="inttype", init="_PrimitiveLoader(int)", instance_type="int")
+_float_type_def: Final = TypeDef(
+    name="floattype", init="_PrimitiveLoader(float)", instance_type="float"
+)
+_bool_type_def: Final = TypeDef(
+    name="booltype", init="_PrimitiveLoader(bool)", instance_type="bool"
+)
+_null_type_def: Final = TypeDef(
+    name="None_type", init="_PrimitiveLoader(type(None))", instance_type="None"
+)
+_any_type_def: Final = TypeDef(name="Any_type", init="_AnyLoader()", instance_type="Any")
 
 prims: Final = {
     "http://www.w3.org/2001/XMLSchema#string": _string_type_def,
@@ -59,7 +65,7 @@ def fmt(text: str, indent: int) -> str:
         black.format_str(
             text,
             mode=black.mode.Mode(
-                target_versions={black.mode.TargetVersion.PY36}, line_length=88 - indent
+                target_versions={black.mode.TargetVersion.PY310}, line_length=88 - indent
             ),
         ),
         " " * indent,
@@ -145,7 +151,9 @@ def begin_class(
         optional_fields: set[str],
     ) -> None:
         classname = self.safe_name(classname)
-        self.current_class_is_abstract = abstract
+        self.current_optional_fields = optional_fields
+        self.current_fieldtypes: dict[str, TypeDef] = {}
+        self.current_class_is_abstract: bool = abstract
 
         if extends:
             ext = ", ".join(self.safe_name(e) for e in extends)
@@ -169,65 +177,17 @@ def begin_class(
             self.out.write("    pass\n\n\n")
             return
 
-        idfield_safe_name: Final = self.safe_name(idfield) if idfield != "" else None
-        if idfield_safe_name is not None:
-            self.out.write(f"    {idfield_safe_name}: str\n\n")
+        self.current_idfield: str = self.safe_name(idfield) if idfield != "" else None
+        if self.current_idfield is not None:
+            self.out.write(f"    {self.current_idfield}: str\n\n")
 
-        required_field_names: Final = [f for f in field_names if f not in optional_fields]
-        optional_field_names: Final = [f for f in field_names if f in optional_fields]
-
-        safe_inits: Final[list[str]] = ["        self,"]
-        safe_inits.extend(
-            [f"        {self.safe_name(f)}: Any," for f in required_field_names if f != "class"]
-        )
-        safe_inits.extend(
-            [
-                f"        {self.safe_name(f)}: Any | None = None,"
-                for f in optional_field_names
-                if f != "class"
-            ]
-        )
-        self.out.write(
-            "    def __init__(\n"
-            + "\n".join(safe_inits)
-            + "\n        extension_fields: MutableMapping[str, Any] | None = None,"
-            + "\n        loadingOptions: LoadingOptions | None = None,"
-            + "\n    ) -> None:\n"
-            + """        if extension_fields:
-            self.extension_fields = extension_fields
-        else:
-            self.extension_fields = CommentedMap()
-        if loadingOptions:
-            self.loadingOptions = loadingOptions
-        else:
-            self.loadingOptions = LoadingOptions()
-"""
-        )
-        field_inits = ""
-        for name in field_names:
-            if name == "class":
-                field_inits += """        self.class_: Final[str] = "{}"
-""".format(
-                    classname
-                )
-            elif name == idfield_safe_name:
-                field_inits += """        self.{0} = {0} if {0} is not None else "_:" + str(_uuid__.uuid4())
-""".format(
-                    self.safe_name(name)
-                )
-            else:
-                field_inits += """        self.{0} = {0}
-""".format(
-                    self.safe_name(name)
-                )
         field_eqs: Final = []
         field_hashes: Final = []
         for name in field_names:
             field_eqs.append("self.{0} == other.{0}".format(self.safe_name(name)))
             field_hashes.append(f"self.{self.safe_name(name)}")
         field_eq: Final = " and\n                    ".join(field_eqs)
         field_hash: Final = ",\n            ".join(field_hashes)
-        self.out.write(field_inits)
         self.out.write(
             "\n"
             + fmt(
@@ -341,7 +301,65 @@ def end_class(self, classname: str, field_names: list[str]) -> None:
 """
         )
 
-        self.serializer.write("        return r\n\n")
+        self.serializer.write("        return r\n")
+
+        required_field_names: Final = [
+            f for f in field_names if f not in self.current_optional_fields
+        ]
+        optional_field_names: Final = [f for f in field_names if f in self.current_optional_fields]
+
+        idfield_safe_name: Final = (
+            self.safe_name(self.current_idfield) if self.current_idfield != None else None
+        )
+        safe_inits: Final[list[str]] = ["        self,"]
+        safe_inits.extend(
+            [
+                f"        {self.safe_name(f)}: {self.current_fieldtypes[self.safe_name(f)].instance_type},"
+                for f in required_field_names
+                if f != "class"
+            ]
+        )
+        safe_inits.extend(
+            [
+                f"        {self.safe_name(f)}: {self.current_fieldtypes[self.safe_name(f)].instance_type} = None,"
+                for f in optional_field_names
+                if f != "class"
+            ]
+        )
+        self.serializer.write(
+            "\n    def __init__(\n"
+            + "\n".join(safe_inits)
+            + "\n        extension_fields: MutableMapping[str, Any] | None = None,"
+            + "\n        loadingOptions: LoadingOptions | None = None,"
+            + "\n    ) -> None:\n"
+            + """        if extension_fields:
+            self.extension_fields = extension_fields
+        else:
+            self.extension_fields = CommentedMap()
+        if loadingOptions:
+            self.loadingOptions = loadingOptions
+        else:
+            self.loadingOptions = LoadingOptions()
+"""
+        )
+        field_inits = ""
+        for name in field_names:
+            if name == "class":
+                field_inits += """        self.class_: Final[str] = "{}"
+""".format(
+                    classname
+                )
+            elif name == idfield_safe_name:
+                field_inits += """        self.{0} = {0} if {0} is not None else "_:" + str(_uuid__.uuid4())
+""".format(
+                    self.safe_name(name)
+                )
+            else:
+                field_inits += """        self.{0} = {0}
+""".format(
+                    self.safe_name(name)
+                )
+        self.serializer.write(f"{field_inits}\n")
 
         self.serializer.write(
             fmt(
@@ -384,40 +402,46 @@ def type_loader(
         """Parse the given type declaration and declare its components."""
         match type_declaration:
             case MutableSequence():
-                sub_names1: Final = list(
-                    dict.fromkeys([self.type_loader(i).name for i in type_declaration])
-                )
+                sub_types1: Final = [self.type_loader(i) for i in type_declaration]
+                sub_names1: Final = [t.name for t in sub_types1]
                 return self.declare_type(
                     TypeDef(
-                        "union_of_{}".format("_or_".join(sub_names1)),
-                        "_UnionLoader(({},))".format(", ".join(sub_names1)),
+                        name="union_of_{}".format("_or_".join(sub_names1)),
+                        init="_UnionLoader(({},))".format(", ".join(sub_names1)),
+                        instance_type=" | ".join({t.instance_type or "" for t in sub_types1}),
                     )
                 )
             case {"type": "array" | "https://w3id.org/cwl/salad#array", "items": items}:
                 i1: Final = self.type_loader(items)
                 return self.declare_type(
                     TypeDef(
-                        f"array_of_{i1.name}",
-                        f"_ArrayLoader({i1.name})",
+                        name=f"array_of_{i1.name}",
+                        init=f"_ArrayLoader({i1.name})",
+                        instance_type=f"MutableSequence[{i1.instance_type}]",
                     )
                 )
             case {"type": "map" | "https://w3id.org/cwl/salad#map", "values": values, **rest}:
                 i2: Final = self.type_loader(values)
                 name = self.safe_name(str(rest["name"])) if "name" in rest else None
                 anon_type = self.declare_type(
                     TypeDef(
-                        f"map_of_{i2.name}",
-                        "_MapLoader({}, {}, {}, {})".format(
+                        name=f"map_of_{i2.name}",
+                        init="_MapLoader({}, {}, {}, {})".format(
                             i2.name,
                             f"'{name}'",  # noqa: B907
                             f"'{container}'" if container is not None else None,  # noqa: B907
                             no_link_check,
                         ),
+                        instance_type=f"MutableMapping[str, Any]",
                     )
                 )
                 if "name" in rest:
                     return self.declare_type(
-                        TypeDef(self.safe_name(str(rest["name"])) + "Loader", anon_type.name)
+                        TypeDef(
+                            name=self.safe_name(str(rest["name"])) + "Loader",
+                            init=anon_type.name,
+                            instance_type=anon_type.instance_type,
+                        )
                     )
                 else:
                     return anon_type
@@ -438,26 +462,30 @@ def type_loader(
                     docstring = f'\n"""\n{formated_doc}\n"""'
                 else:
                     docstring = ""
+                sym_names: Final = [schema.avro_field_name(sym) for sym in symbols]
+                sym_literals: Final = [f'"{s}"' for s in sym_names]
                 return self.declare_type(
                     TypeDef(
-                        self.safe_name(name) + "Loader",
-                        '_EnumLoader(("{}",), "{}"){}'.format(
-                            '", "'.join(schema.avro_field_name(sym) for sym in symbols),
+                        name=self.safe_name(name) + "Loader",
+                        init='_EnumLoader(("{}",), "{}"){}'.format(
+                            '", "'.join(sym_names),
                             self.safe_name(name),
                             docstring,
                         ),
+                        instance_type=f"Literal[{', '.join(sym_literals)}]",
                     )
                 )
 
             case {"type": "record" | "https://w3id.org/cwl/salad#record", "name": name, **rest}:
                 return self.declare_type(
                     TypeDef(
-                        self.safe_name(name) + "Loader",
-                        "_RecordLoader({}, {}, {})".format(
+                        name=self.safe_name(name) + "Loader",
+                        init="_RecordLoader({}, {}, {})".format(
                             self.safe_name(name),
                             f"'{container}'" if container is not None else None,  # noqa: B907
                             no_link_check,
                         ),
+                        instance_type=self.safe_name(name),
                         abstract=bool(rest.get("abstract", False)),
                     )
                 )
@@ -469,15 +497,23 @@ def type_loader(
             }:
                 # Declare the named loader to handle recursive union definitions
                 loader_name = self.safe_name(name) + "Loader"
-                loader_type = TypeDef(loader_name, f"_UnionLoader((), '{loader_name}')")
+                loader_type = TypeDef(
+                    name=loader_name,
+                    init=f"_UnionLoader((), '{loader_name}')",
+                    instance_type=self.safe_name(name),
+                )
                 self.declare_type(loader_type)
                 # Parse inner types
-                sub_names2: Final = list(dict.fromkeys([self.type_loader(i).name for i in names]))
+                sub_types2: Final = [self.type_loader(i) for i in names]
+                sub_names2: Final = list(dict.fromkeys(t.name for t in sub_types2))
                 # Register lazy initialization for the loader
                 self.add_lazy_init(
                     LazyInitDef(
-                        loader_name,
-                        "{}.add_loaders(({},))".format(loader_name, ", ".join(sub_names2)),
+                        name=loader_name,
+                        init="{}.add_loaders(({},))".format(loader_name, ", ".join(sub_names2)),
+                        instance_type=f'{self.safe_name(name)}: TypeAlias = "'
+                        + " | ".join({t.instance_type for t in sub_types2})
+                        + '"',
                     )
                 )
                 return loader_type
@@ -490,8 +526,9 @@ def type_loader(
             case "Expression" | "https://w3id.org/cwl/cwl#Expression" as decl:
                 return self.declare_type(
                     TypeDef(
-                        self.safe_name(decl) + "Loader",
-                        "_ExpressionLoader(str)",
+                        name=self.safe_name(decl) + "Loader",
+                        init="_ExpressionLoader(str)",
+                        instance_type="str",
                     )
                 )
             case str(decl):
@@ -545,6 +582,7 @@ def declare_field(
     ) -> None:
         if self.current_class_is_abstract:
             return
+        self.current_fieldtypes[self.safe_name(name)] = fieldtype
 
         if optional:
             self.out.write(f"""        {self.safe_name(name)} = None\n""")
@@ -728,11 +766,12 @@ def uri_loader(
         """Construct the TypeDef for the given URI loader."""
         return self.declare_type(
             TypeDef(
-                f"uri_{inner.name}_{scoped_id}_{vocab_term}_{ref_scope}_{no_link_check}",
-                f"_URILoader({inner.name}, {scoped_id}, {vocab_term}, {ref_scope}, {no_link_check})",
+                name=f"uri_{inner.name}_{scoped_id}_{vocab_term}_{ref_scope}_{no_link_check}",
+                init=f"_URILoader({inner.name}, {scoped_id}, {vocab_term}, {ref_scope}, {no_link_check})",
                 is_uri=True,
                 scoped_id=scoped_id,
                 ref_scope=ref_scope,
+                instance_type=inner.instance_type,
             )
         )
 
@@ -742,27 +781,30 @@ def idmap_loader(
         """Construct the TypeDef for the given mapped ID loader."""
         return self.declare_type(
             TypeDef(
-                f"idmap_{self.safe_name(field)}_{inner.name}",
-                f"_IdMapLoader({inner.name}, '{map_subject}', '{map_predicate}')",  # noqa: B907
+                name=f"idmap_{self.safe_name(field)}_{inner.name}",
+                init=f"_IdMapLoader({inner.name}, '{map_subject}', '{map_predicate}')",  # noqa: B907
+                instance_type=inner.instance_type,
             )
         )
 
     def typedsl_loader(self, inner: TypeDef, ref_scope: int | None) -> TypeDef:
         """Construct the TypeDef for the given DSL loader."""
         return self.declare_type(
             TypeDef(
-                f"typedsl_{self.safe_name(inner.name)}_{ref_scope}",
-                f"_TypeDSLLoader({self.safe_name(inner.name)}, {ref_scope}, "  # noqa: B907
+                name=f"typedsl_{self.safe_name(inner.name)}_{ref_scope}",
+                init=f"_TypeDSLLoader({self.safe_name(inner.name)}, {ref_scope}, "  # noqa: B907
                 f"'{self.salad_version}')",  # noqa: B907
+                instance_type=inner.instance_type,
             )
         )
 
     def secondaryfilesdsl_loader(self, inner: TypeDef) -> TypeDef:
         """Construct the TypeDef for secondary files."""
         return self.declare_type(
             TypeDef(
-                f"secondaryfilesdsl_{inner.name}",
-                f"_UnionLoader((_SecondaryDSLLoader({inner.name}), {inner.name},))",
+                name=f"secondaryfilesdsl_{inner.name}",
+                init=f"_UnionLoader((_SecondaryDSLLoader({inner.name}), {inner.name},))",
+                instance_type=inner.instance_type,
             )
         )
 
@@ -787,6 +829,7 @@ def epilogue(self, root_loader: TypeDef) -> None:
         if self.lazy_inits:
             for lazy_init in self.lazy_inits.values():
                 self.out.write(fmt(f"{lazy_init.init}\n", 0))
+                self.out.write(fmt(f"{lazy_init.instance_type}", 0))
             self.out.write("\n")
 
         self.out.write(

schema_salad/python_codegen_support.py

@@ -17,7 +17,7 @@
 from itertools import chain
 from mypy_extensions import trait
 from typing import Any, Final, Generic, TypeAlias, TypeVar, cast
-from typing import ClassVar  # pylint: disable=unused-import # noqa: F401
+from typing import ClassVar, Literal  # pylint: disable=unused-import # noqa: F401
 from urllib.parse import quote, urldefrag, urlparse, urlsplit, urlunsplit
 from urllib.request import pathname2url