feat: implement IRI access via index operator in LinkedBaseModel

Author: simontaurus

src/oold/model/__init__.py

@@ -1,10 +1,10 @@
 import json
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, overload
 
 import pydantic
 from pydantic import BaseModel
-from typing_extensions import Self  # to support python < 3.11
+from typing_extensions import Self
 
 from oold.model.static import GenericLinkedBaseModel, export_jsonld, import_jsonld
 
@@ -60,18 +60,30 @@ def get_resolver(param: GetResolverParam) -> GetResolverResult:
 class LinkedBaseModelMetaClass(pydantic.main._model_construction.ModelMetaclass):
     def __new__(mcs, name, bases, namespace):
         cls = super().__new__(mcs, name, bases, namespace)
-        schema = {}
-
-        # pydantic v2
-        if "model_config" in namespace:
-            if "json_schema_extra" in namespace["model_config"]:
-                schema = namespace["model_config"]["json_schema_extra"]
 
-        if "iri" in schema:
-            iri = schema["iri"]
-            _types[iri] = cls
+        if hasattr(cls, "get_cls_iri"):
+            iri = cls.get_cls_iri()
+            if iri is not None:
+                _types[iri] = cls
         return cls
 
+    # override operators, see https://docs.python.org/3/library/operator.html
+
+    @overload
+    def __getitem__(cls: "LinkedBaseModel", item: str) -> Self:
+        ...
+
+    @overload
+    def __getitem__(cls: "LinkedBaseModel", item: List[str]) -> List[Self]:
+        ...
+
+    def __getitem__(
+        cls: "LinkedBaseModel", item: Union[str, List[str]]
+    ) -> Union[Self, List[Self]]:
+        """Allow access to the class by its IRI."""
+        result = cls._resolve(item if isinstance(item, list) else [item])
+        return result[item] if isinstance(item, str) else [result[i] for i in item]
+
 
 # the following switch ensures that autocomplete works in IDEs like VSCode
 if TYPE_CHECKING:
@@ -92,6 +104,20 @@ class LinkedBaseModel(_LinkedBaseModel):
 
     __iris__: Optional[Dict[str, Union[str, List[str]]]] = {}
 
+    @classmethod
+    def get_cls_iri(cls) -> str:
+        """Return the unique IRI of the class.
+        Overwrite this method in the subclass."""
+        schema = {}
+        if hasattr(cls, "__config__"):
+            if hasattr(cls.__config__, "schema_extra"):
+                schema = cls.__config__.schema_extra
+
+        if "iri" in schema:
+            return schema["iri"]
+        else:
+            return None
+
     def get_iri(self) -> str:
         """Return the unique IRI of the object.
         Overwrite this method in the subclass."""

src/oold/model/v1/__init__.py

@@ -1,9 +1,10 @@
 import json
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union, overload
 
 import pydantic
 from pydantic.v1 import BaseModel, PrivateAttr
+from typing_extensions import Self
 
 from oold.model.static import GenericLinkedBaseModel, export_jsonld, import_jsonld
 
@@ -62,18 +63,30 @@ def get_resolver(param: GetResolverParam) -> GetResolverResult:
 class LinkedBaseModelMetaClass(pydantic.v1.main.ModelMetaclass):
     def __new__(mcs, name, bases, namespace):
         cls = super().__new__(mcs, name, bases, namespace)
-        schema = {}
-
-        # pydantic v1
-        if "Config" in namespace:
-            if "schema_extra" in namespace["Config"].__dict__:
-                schema = namespace["Config"].schema_extra
 
-        if "iri" in schema:
-            iri = schema["iri"]
-            _types[iri] = cls
+        if hasattr(cls, "get_cls_iri"):
+            iri = cls.get_cls_iri()
+            if iri is not None:
+                _types[iri] = cls
         return cls
 
+    # override operators, see https://docs.python.org/3/library/operator.html
+
+    @overload
+    def __getitem__(cls: "LinkedBaseModel", item: str) -> Self:
+        ...
+
+    @overload
+    def __getitem__(cls: "LinkedBaseModel", item: List[str]) -> List[Self]:
+        ...
+
+    def __getitem__(
+        cls: "LinkedBaseModel", item: Union[str, List[str]]
+    ) -> Union[Self, List[Self]]:
+        """Allow access to the class by its IRI."""
+        result = cls._resolve(item if isinstance(item, list) else [item])
+        return result[item] if isinstance(item, str) else [result[i] for i in item]
+
 
 # the following switch ensures that autocomplete works in IDEs like VSCode
 if TYPE_CHECKING:
@@ -94,6 +107,20 @@ class LinkedBaseModel(_LinkedBaseModel):
 
     __iris__: Optional[Dict[str, Union[str, List[str]]]] = PrivateAttr()
 
+    @classmethod
+    def get_cls_iri(cls) -> str:
+        """Return the unique IRI of the class.
+        Overwrite this method in the subclass."""
+        schema = {}
+        if hasattr(cls, "__config__"):
+            if hasattr(cls.__config__, "schema_extra"):
+                schema = cls.__config__.schema_extra
+
+        if "iri" in schema:
+            return schema["iri"]
+        else:
+            return None
+
     def get_iri(self) -> str:
         """Return the unique IRI of the object.
         Overwrite this method in the subclass."""

tests/test_oold.py

@@ -217,6 +217,12 @@ def export_json(obj):
     else:
         assert False, "ValueError not raised"
 
+    # test index operator for getting objects by IRI
+    f = model.Foo["ex:f"]
+    assert f.id == "ex:f"
+    [b1, b2] = model.Bar[["ex:b1", "ex:b2"]]
+    assert b1.id == "ex:b1" and b2.id == "ex:b2"
+
 
 def test_core():
     _run(pydantic_version="v1")