feat: build typed kwargs dicts

src/ome_autogen/main.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
 import os
+import re
 import subprocess
 import sys
 from pathlib import Path
 from shutil import rmtree
+from typing import Any
 
 from ome_autogen import _util
-from ome_autogen._config import get_config
+from ome_autogen._config import PYDANTIC_SUPPORT, get_config
 from ome_autogen._transformer import OMETransformer
 
 BLACK_LINE_LENGTH = 88
@@ -49,6 +51,7 @@ def build_model(
         _print_gray("Writing Files...")
         transformer.process_classes()
 
+    _build_typed_dicts(package_dir)
     if do_formatting:
         _fix_formatting(package_dir, ruff_ignore)
 
@@ -99,3 +102,89 @@ def _print_green(text: str) -> None:
         # UnicodeEncodeError: 'charmap' codec can't encode character '\u2713'
         text = f"\033[92m\033[1m{text}\033[0m"
     print(text)
+
+
+KWARGS_MODULE = """
+from __future__ import annotations
+from typing_extensions import TypedDict, TypeAlias
+from typing import Union, List
+from datetime import datetime
+import ome_types.model as ome
+
+class RefDict(TypedDict):
+    id: str
+"""
+
+
+def _build_typed_dicts(package_dir: str) -> None:
+    """Create a TypedDict class for each OMEType subclass.
+
+    Useful for passing kwargs to the constructors.
+
+    https://peps.python.org/pep-0692/
+
+    def foo(**kwargs: Unpack[ome.ImageDict]) -> None:
+        ...
+    """
+    # sourcery skip: assign-if-exp, reintroduce-else
+    try:
+        from pydantic._internal._repr import display_as_type
+    except ImportError:
+        # don't try to do this on pydantic1
+        return
+    if PYDANTIC_SUPPORT == "v1":
+        return
+
+    from ome_types import model
+    from ome_types._mixins._base_type import OMEType
+
+    ome_models = {
+        name: obj
+        for name, obj in vars(model).items()
+        if isinstance(obj, type) and issubclass(obj, OMEType) and obj.__annotations__
+    }
+
+    def _disp_type(obj: Any) -> str:
+        x = display_as_type(obj).replace("NoneType", "None")
+        if "ForwardRef" in x:
+            #  replace "List[ForwardRef('Map.M')]" with "List[Map.M]"
+            x = re.sub(r"ForwardRef\('([a-zA-Z_.]*)'\)", r"\1", x)
+        x = re.sub("ome_types\._autogenerated\.ome_2016_06.[^.]+.", "", x)
+        return x
+
+    # add TypedDicts for all models
+    module = KWARGS_MODULE
+    SUFFIX = "Dict"
+    CLASS = "class {name}(TypedDict, total=False):\n\t{fields}\n\n"
+    for cls_name, m in sorted(ome_models.items()):
+        if cls_name.endswith("Ref"):
+            module += f"{cls_name}: TypeAlias = RefDict\n"
+        else:
+            _fields = [
+                f"{k}: {_disp_type(v.annotation)}"
+                for k, v in sorted(m.model_fields.items())
+            ]
+            if _fields:
+                module += CLASS.format(
+                    name=f"{m.__name__}{SUFFIX}", fields="\n\t".join(_fields)
+                )
+            else:
+                module += (
+                    f"class {m.__name__}{SUFFIX}(TypedDict, total=False):\n\tpass\n\n"
+                )
+
+    # fix name spaces
+    # prefix all remaining capitalized words with ome.
+    def _repl(match: re.Match) -> str:
+        word = match[1]
+        if word in {"None", "True", "False", "Union", "List", "TypedDict", "TypeAlias"}:
+            return word
+        if word.endswith(SUFFIX):
+            return word
+        if word in ome_models:
+            return word + SUFFIX
+        # the rest are enums, they can be passed as strings
+        return f"ome.{word} | str"
+
+    module = re.sub(r"\b([A-Z][a-zA-Z_^.]*)\b", _repl, module)
+    (Path(package_dir) / "kwargs.py").write_text(module)

src/ome_types/model/__init__.py

@@ -22,6 +22,8 @@
     from importlib.machinery import ModuleSpec
     from types import ModuleType
 
+    from ome_types._autogenerated.ome_2016_06 import kwargs as kwargs
+
 # ---------------------------------------------------------------------
 # Below here is logic to allow importing from ome_types._autogenerated.ome_2016_06
 # from ome_types.model.* (to preserve backwards compatibility)