feat: allow model registration using directory path

Author: matmoncon

.pylintrc

@@ -7,6 +7,7 @@ unsafe-load-any-extension=no
 
 [MESSAGES CONTROL]
 disable=abstract-method,
+        inconsistent-quotes,
         protected-access,
         broad-exception-caught,
         ungrouped-imports,

pyneo4j_ogm/core/client.py

@@ -2,6 +2,8 @@
 Pyneo4j database client class for running operations on the database.
 """
 
+import importlib.util
+import inspect
 import os
 from enum import Enum
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union, cast
@@ -14,7 +16,6 @@
 from pyneo4j_ogm.core.node import NodeModel
 from pyneo4j_ogm.core.relationship import RelationshipModel
 from pyneo4j_ogm.exceptions import (
-    AlreadyRegistered,
     InvalidBookmark,
     InvalidEntityType,
     InvalidLabelOrType,
@@ -23,7 +24,6 @@
     TransactionInProgress,
     UnsupportedNeo4jVersion,
 )
-from pyneo4j_ogm.fields.settings import NodeModelSettings, RelationshipModelSettings
 from pyneo4j_ogm.logger import logger
 from pyneo4j_ogm.pydantic_utils import get_field_type, get_model_fields
 from pyneo4j_ogm.queries.query_builder import QueryBuilder
@@ -138,6 +138,41 @@ async def connect(
         logger.info("Connected to database")
         return self
 
+    @ensure_connection
+    async def register_models_dir(self, dir_path: str) -> None:
+        """
+        Registers all models in a directory and all subdirectories.
+        """
+        logger.info("Registering models in directory %s", dir_path)
+        for root, _, files in os.walk(dir_path):
+            # Check all files for models
+            logger.debug("Checking %s files for models", len(files))
+            for file in files:
+                if not file.endswith(".py"):
+                    continue
+
+                filepath = os.path.join(root, file)
+
+                # Load the module
+                logger.debug("Found file %s, importing", filepath)
+                module_name = os.path.splitext(os.path.basename(filepath))[0]
+                spec = importlib.util.spec_from_file_location(module_name, filepath)
+
+                if spec is None or spec.loader is None:
+                    raise ImportError(f"Could not import migration file {filepath}")
+
+                module = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(module)
+
+                for member in inspect.getmembers(
+                    module,
+                    lambda x: inspect.isclass(x)
+                    and issubclass(x, (NodeModel, RelationshipModel))
+                    and x is not NodeModel
+                    and x is not RelationshipModel,
+                ):
+                    self.models.add(member[1])
+
     @ensure_connection
     async def register_models(self, models: List[Type[Union[NodeModel, RelationshipModel]]]) -> None:
         """
@@ -153,23 +188,6 @@ async def register_models(self, models: List[Type[Union[NodeModel, RelationshipM
         logger.info("Registering models %s with client %s", models, self)
 
         for model in models:
-            for registered_model in self.models:
-                registered_model_settings = registered_model.model_settings()
-                model_settings = model.model_settings()
-
-                if (
-                    isinstance(registered_model_settings, RelationshipModelSettings)
-                    and isinstance(model_settings, RelationshipModelSettings)
-                    and registered_model_settings.type == model_settings.type
-                ):
-                    raise AlreadyRegistered(cast(str, model_settings.type))
-                elif (
-                    isinstance(registered_model_settings, NodeModelSettings)
-                    and isinstance(model_settings, NodeModelSettings)
-                    and set(registered_model_settings.labels) == set(model_settings.labels)
-                ):
-                    raise AlreadyRegistered(cast(str, model_settings.labels))
-
             if issubclass(model, (NodeModel, RelationshipModel)):
                 logger.debug("Found valid mode %s, registering with client", model.__name__)
 

pyneo4j_ogm/exceptions.py

@@ -215,14 +215,3 @@ class ListItemNotEncodable(Pyneo4jException):
 
     def __init__(self, *args: object) -> None:
         super().__init__("List item is not JSON encodable and can not be stored inside the database", *args)
-
-
-class AlreadyRegistered(Pyneo4jException):
-    """
-    Multiple models are using the same labels/type as a already registered model.
-    """
-
-    def __init__(self, labels_or_type: Union[Set[str], str], *args: object) -> None:
-        received = f"[{', '.join(labels_or_type)}]" if isinstance(labels_or_type, set) else f"[{labels_or_type}]"
-        label_or_type = "labels" if isinstance(labels_or_type, set) else "type"
-        super().__init__(f"A model is using the same {label_or_type} as another model. Got {received}", *args)

tests/core/test_client.py

@@ -14,7 +14,6 @@
 from pyneo4j_ogm.core.node import NodeModel
 from pyneo4j_ogm.core.relationship import RelationshipModel
 from pyneo4j_ogm.exceptions import (
-    AlreadyRegistered,
     InvalidEntityType,
     InvalidLabelOrType,
     MissingDatabaseURI,
@@ -25,6 +24,12 @@
 from pyneo4j_ogm.fields.property_options import WithOptions
 from pyneo4j_ogm.logger import logger
 from tests.fixtures.db_setup import client, session
+from tests.fixtures.models.models_top import ModelOne, ModelTwo
+from tests.fixtures.models.nested.deeply_nested.model_deeply_nested import (
+    ModelFive,
+    ModelSix,
+)
+from tests.fixtures.models.nested.model_nested import ModelFour, ModelThree
 
 
 class CypherResolvingNode(NodeModel):
@@ -41,19 +46,6 @@ class Settings:
         type = "TEST_RELATIONSHIP"
 
 
-async def test_duplicate_model_register(client: Pyneo4jClient):
-    class TestNode(NodeModel):
-        name: str
-
-        class Settings:
-            labels = {"TestNode"}
-
-    await client.register_models([TestNode])
-
-    with pytest.raises(AlreadyRegistered):
-        await client.register_models([TestNode])
-
-
 async def test_batch(client: Pyneo4jClient, session: AsyncSession):
     async with client.batch():
         await client.cypher("CREATE (n:Node) SET n.name = $name", parameters={"name": "TestName"})
@@ -621,3 +613,8 @@ async def test_drop_indexes(client: Pyneo4jClient, session: AsyncSession):
     await query_results.consume()
 
     assert len(results) == 0
+
+
+async def test_register_models_dir(client: Pyneo4jClient):
+    await client.register_models_dir("tests/fixtures/models")
+    assert len(client.models) == 6

tests/fixtures/models/models_top.py

@@ -0,0 +1,10 @@
+from pyneo4j_ogm.core.node import NodeModel
+from pyneo4j_ogm.core.relationship import RelationshipModel
+
+
+class ModelOne(NodeModel):
+    pass
+
+
+class ModelTwo(RelationshipModel):
+    pass

tests/fixtures/models/nested/deeply_nested/model_deeply_nested.py

@@ -0,0 +1,10 @@
+from pyneo4j_ogm.core.node import NodeModel
+from pyneo4j_ogm.core.relationship import RelationshipModel
+
+
+class ModelFive(NodeModel):
+    pass
+
+
+class ModelSix(RelationshipModel):
+    pass

tests/fixtures/models/nested/deeply_nested/other_classes.py

@@ -0,0 +1,2 @@
+class NotMePleaseButDeeplyNested:
+    pass

tests/fixtures/models/nested/model_nested.py

@@ -0,0 +1,10 @@
+from pyneo4j_ogm.core.node import NodeModel
+from pyneo4j_ogm.core.relationship import RelationshipModel
+
+
+class ModelThree(NodeModel):
+    pass
+
+
+class ModelFour(RelationshipModel):
+    pass

tests/fixtures/models/nested/no_models.py

@@ -0,0 +1,2 @@
+class NotMePlease:
+    pass