Fix RootModels for Pydantic 2.10

CHANGELOG.md

@@ -12,7 +12,22 @@ While the project is still on major version 0, breaking changes may be introduce
 
 <!-- changelog follows -->
 
-<!-- ## Unreleased -->
+## Unreleased
+
+### Changed
+
+- Models updated to API schema from [c97253f](https://github.com/goharbor/harbor/blob/4a12623459a754ff4d07fbd1cddb4df436e8524c/api/v2.0/swagger.yaml)
+- All root models now share a common `harborapi.models.base.RootModel` base class.
+
+### Fixed
+
+- Root models failing to build on Pydantic 2.10.
+
+### Removed
+
+- Generic aliased root models:
+  - `harborapi.models.base.StrDictRootModel`
+  - `harborapi.models.base.StrRootModel`
 
 ## [0.25.3](https://github.com/unioslo/harborapi/tree/harborapi-v0.25.3) - 2024-08-26
 

codegen/ast/fragments/main/registryproviders.py

@@ -8,7 +8,7 @@
 from pydantic import Field
 
 
-class RegistryProviders(RootModel):
+class RegistryProviders(RootModel[Dict[str, RegistryProviderInfo]]):
     root: Dict[str, RegistryProviderInfo] = Field(
         default={},
         description="The registry providers. Each key is the name of the registry provider.",

codegen/ast/parser.py

@@ -310,7 +310,7 @@ def modify_module(tree: ast.Module, fragment_dir: FragmentDir) -> ast.Module:
 
 # Imports that should be added to every file
 ADD_IMPORTS = {
-    "harborapi.models.base": ["StrDictRootModel", "StrRootModel"],
+    "harborapi.models.base": [],
 }  # type: dict[str, list[str]] # module: list[import_name]
 
 
@@ -321,16 +321,21 @@ def add_imports(tree: ast.Module) -> ast.Module:
         if isinstance(node, ast.ImportFrom):
             if node.module in ADD_IMPORTS:
                 node_names = [node.name for node in node.names]
+                if not node.module:
+                    continue
                 for name in ADD_IMPORTS[node.module]:
                     if name not in node_names:
                         node.names.append(ast.alias(name=name, asname=None))
                 added.add(node.module)
     # Remaining imports that were not appended to existing imports
     for name in set(ADD_IMPORTS) - added:
         names = [ast.alias(name=name, asname=None) for name in ADD_IMPORTS[name]]
+        if not names:
+            continue  # no imports to add from module
+
+        # Inserts `from module import name1, name2, ...`
         tree.body.insert(1, ast.ImportFrom(module=name, names=names, level=0))
         # Assume from __future__ is at the top of the file
-        # Regardless, we automatically sort imports afterwards anyway
     return tree
 
 
@@ -365,107 +370,6 @@ class Foo(RootModel): # lacks parametrized base
     raise ValueError(f"Class definition '{classdef.name}' does not have a root field.")
 
 
-def fix_rootmodel_base(classdef: ast.ClassDef) -> None:
-    """Adds the appropriate subclass as the base of a RootModel type.
-
-    Depending on the root value annotation, the function will assign one of two
-    bases:
-
-    - `StrDictRootModel` if the root value annotation is `Optional[Dict[str, T]]`
-    - `StrRootModel` if the root value annotation is `str`
-
-    As of goharbor/harbor@5c02fd8, there are no models encapsulating dicts
-    whose root value type is `Dict[str, T]`; they are always `Optional[Dict[str, T]]`.
-
-    Examples
-    --------
-
-    ```
-    class Foo(RootModel):
-        root: Optional[Dict[str, str]]
-    # ->
-    class Foo(StrDictRootModel[str]):
-        root: Optional[Dict[str, str]]
-    ```
-
-    Also works for str root models:
-    ```
-    class Bar(RootModel):
-        root: str
-    # ->
-    class Bar(StrRootModel):
-        root: str
-    ```
-
-    See also
-    --------
-    `harborapi.models.base.StrRootModel`
-    `harborapi.models.base.StrDictRootModel`
-    """
-    # Determine what sort of root model we are dealing with
-    root_type = get_rootmodel_type(classdef)
-    base = "RootModel"
-    vt = "Any"
-    # Root type is a string annotation
-    # e.g. root: "Dict[str, str]"
-    if isinstance(root_type, ast.Name):
-        # HACK: this will break for root models with more complicated signatures,
-        # but we are not dealing with that right now
-        if "Dict[str" in root_type.id:
-            base = "StrDictRootModel"
-            # HACK: create Python statement with the type annotation
-            # and then parse it to get the AST
-            # Say our annotation is `Dict[str, str]`, we want to pass
-            # `str` as the type parameter to `StrDictRootModel`.
-            annotation = ast.parse(f"var: {root_type.id}").body[0].annotation
-            # If the annotation is Optional[Dict[str, str]], then we need
-            # to go through one more slice to get the value type
-            # i.e. Optional[Dict[str, str]] -> Dict[str, str] -> str
-            if "Optional" in root_type.id:
-                slc = annotation.slice.slice
-            else:
-                slc = annotation.slice
-            vt = slc.elts[1].id  # (KT, VT)
-        elif root_type.id == "str":
-            base = "StrRootModel"
-    # Root type is an annotation with a subscript, e.g. Dict[str, T]
-    # or Optional[Dict[str, T]]
-    elif isinstance(root_type, ast.Subscript):
-        # Inspect the AST to determine the type of root model
-        # If annotation is wrapped in Optional[], we need to get the inner slice
-        if getattr(root_type.value, "id", None) == "Optional":
-            inner_root_type = getattr(root_type, "slice")
-        else:
-            inner_root_type = root_type
-        if getattr(inner_root_type.value, "id", None) == "Dict":
-            base = "StrDictRootModel"
-            vt = inner_root_type.slice.elts[1].id  # (KT, VT)
-        # TODO: handle list root types
-    else:
-        raise ValueError(f"Invalid root type: {root_type}")
-
-    # Construct the node for the class's new base
-    if base == "StrDictRootModel":
-        classdef.bases = [
-            ast.Subscript(
-                value=ast.Name(id="StrDictRootModel"),
-                slice=ast.Index(ast.Name(id=vt)),
-            )
-        ]
-    else:
-        # Otherwise, we use the base we determined earlier
-        classdef.bases = [ast.Name(id=base)]
-
-
-def fix_rootmodels(tree: ast.Module, classdefs: dict[str, ast.ClassDef]) -> ast.Module:
-    for node in ast.walk(tree):
-        if not isinstance(node, ast.ClassDef):
-            continue
-        if _get_class_base_name(node) == "RootModel":
-            fix_rootmodel_base(node)
-    return tree
-
-
 def insert_or_update_classdefs(
     tree: ast.Module, classdefs: dict[str, ast.ClassDef]
 ) -> ast.Module:
@@ -554,7 +458,6 @@ def add_fragments(tree: ast.Module, directory: Path) -> ast.Module:
             statements["stmts"].extend(stmts["stmts"])
     new_tree = insert_or_update_classdefs(tree, classdefs)
     new_tree = insert_statements(new_tree, statements)
-    new_tree = fix_rootmodels(new_tree, classdefs)
     return new_tree
 
 

harborapi/models/base.py

@@ -9,10 +9,13 @@
 from __future__ import annotations
 
 from typing import Any
-from typing import Dict
+from typing import Generator
 from typing import Iterable
+from typing import Mapping
 from typing import Optional
+from typing import Sequence
 from typing import Set
+from typing import Tuple
 from typing import Type
 from typing import TypeVar
 
@@ -54,40 +57,33 @@
 class RootModel(PydanticRootModel[T]):
     model_config = ConfigDict(validate_assignment=True)
 
+    root: T
+
     def __bool__(self) -> bool:
         return bool(self.root)
 
-
-class StrDictRootModel(RootModel[Optional[Dict[str, T]]]):
-    # All JSON keys are string, so the key type does need to be
-    # parameterized with a generic type.
-
-    def __iter__(self) -> Any:
+    def __iter__(self) -> Generator[Tuple[Any, Any], None, None]:
         # TODO: fix API spec so root  types can never be none, only
         # the empty container. That way we can always iterate and access
         # without checking for None.
-        if self.root is not None:
-            return iter(self.root)
-        return iter([])
+        if isinstance(self.root, Iterable):
+            yield from iter(self.root)  # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType]
+        else:
+            yield from iter([])
 
-    def __getitem__(self, item: str) -> Optional[T]:
-        if self.root is not None:
-            return self.root[item]
+    def __getitem__(self, key: Any) -> Any:
+        if isinstance(self.root, (Mapping, Sequence)):
+            return self.root[key]  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
         return None
 
     # Enables dot access to dict keys for backwards compatibility
     def __getattr__(self, attr: str) -> T:
         try:
-            return self.root[attr]  # type: ignore # forego None check and let KeyError raise
-        except (KeyError, TypeError):
+            return self.root[attr]  # pyright: ignore[reportUnknownVariableType, reportIndexIssue]
+        except (KeyError, TypeError, IndexError):
             raise AttributeError(f"{self.__class__.__name__} has no attribute {attr}")
 
 
-class StrRootModel(RootModel[str]):
-    def __str__(self) -> str:
-        return str(self.root)
-
-
 class BaseModel(PydanticBaseModel):
     model_config = ConfigDict(extra="allow", validate_assignment=True, strict=False)