Rework model and instance registry

src/couch_potato/__init__.py

@@ -4,8 +4,8 @@
 from couchbase.cluster import Cluster
 from couchbase.collection import Collection
 
-from src.couch_potato._context import model_create_context
 from src.couch_potato._model import BoundModel, make_meta_model
+from src.couch_potato._registry import get_instance_registry
 from src.couch_potato._types import BucketBind
 
 
@@ -18,19 +18,23 @@ def __init__(self, cluster: Cluster, **kwargs):
         self._binds: Dict[str, BucketBind] = {}
         self._model_binds: Dict[BoundModel, Collection] = {}
 
+        # Add ourselves to the instance registry
+        get_instance_registry().register(self)
+
         # After init, wait until the cluster is ready
         init_timeout = kwargs.get("init_timeout", 5)
         self._cluster.wait_until_ready(timeout=timedelta(seconds=init_timeout))
 
     @property
     def Model(self) -> Type[BoundModel]:
         if self.__model__ is None:
-            with model_create_context():
 
-                class _Model(BoundModel, metaclass=make_meta_model(self)):
-                    @classmethod
-                    def bind(cls) -> Collection:
-                        return self.model_binds[cls]
+            class _Model(
+                BoundModel, metaclass=make_meta_model(self), add_to_registry=False
+            ):
+                @classmethod
+                def bind(cls) -> Collection:
+                    return self.model_binds[cls]
 
             self.__model__ = _Model
         return self.__model__

src/couch_potato/_model.py

@@ -4,7 +4,7 @@
 from couchbase.options import InsertOptions, ReplaceOptions
 from couchbase.scope import Scope
 
-from src.couch_potato._context import MODEL_CREATE_CONTEXT
+from src.couch_potato._registry import get_instance_registry, get_model_registry
 from src.couch_potato._types import BucketBind, Field, ScopeBind
 from src.couch_potato.errors import FieldNotFound, ModelAttributeError, ReadOnlyError
 
@@ -16,7 +16,7 @@
 
 
 class ModelMeta(type):
-    def __new__(cls, name, bases, dct):
+    def __new__(cls, name, bases, dct, add_to_registry: bool = True):
         new_class = super().__new__(cls, name, bases, dct)
 
         # Set up internal attributes on the class defn
@@ -28,24 +28,26 @@ def __new__(cls, name, bases, dct):
                 fields[name] = attr
         setattr(new_class, "__fields__", fields)
 
+        if add_to_registry:
+            get_model_registry().register(new_class)
+
         return new_class
 
 
 def make_meta_model(cp: "CouchPotato"):
     from src.couch_potato.model import KeyGenerator
 
     class BoundModelMeta(ModelMeta):
-        def __new__(cls, name, bases, dct):
-            new_class = super().__new__(cls, name, bases, dct)
+        def __new__(cls, name, bases, dct, add_to_registry: bool = True):
+            new_class = super().__new__(cls, name, bases, dct, add_to_registry=False)
 
             # BaseModel instantiation
-            try:
-                ctx = MODEL_CREATE_CONTEXT.get()
-            except LookupError:
-                pass
-            else:
-                if not ctx.add_to_registry:
-                    return new_class
+            if not add_to_registry:
+                return new_class
+
+            # Add this model to the model registry for the particular instance
+            registry = get_instance_registry()
+            registry.get_model_registry(cp).register(new_class)
 
             # Check validity of the model class
             # 1. Ensure that the model class has a bucket attribute defined

src/couch_potato/_registry.py

@@ -0,0 +1,47 @@
+from typing import TYPE_CHECKING, Dict
+
+from src.couch_potato.errors import RegistryError
+
+if TYPE_CHECKING:
+    from src.couch_potato import CouchPotato
+    from src.couch_potato._model import BaseModel
+
+
+class ModelRegistry:
+    def __init__(self):
+        self._models: Dict[str, "BaseModel"] = {}
+
+    def register(self, model) -> None:
+        if model.__name__ in self._models:
+            raise RegistryError(f"Model with name {model.__name__} already registered")
+        self._models[model.__name__] = model
+
+
+class InstanceRegistry:
+    def __init__(self):
+        self._instances: Dict["CouchPotato", ModelRegistry] = {}
+
+    def register(self, instance) -> None:
+        if instance in self._instances:
+            raise RegistryError("Instance already registered")
+        self._instances[instance] = ModelRegistry()
+
+    def contains(self, instance) -> bool:
+        return instance in self._instances.keys()
+
+    def get_model_registry(self, instance) -> ModelRegistry:
+        if not self.contains(instance):
+            raise RegistryError("Instance is not registered")
+        return self._instances[instance]
+
+
+_model_registry = ModelRegistry()
+_instance_registry = InstanceRegistry()
+
+
+def get_model_registry():
+    return _model_registry
+
+
+def get_instance_registry():
+    return _instance_registry

src/couch_potato/errors.py

@@ -2,6 +2,10 @@ class ModelError(Exception):
     pass
 
 
+class RegistryError(Exception):
+    pass
+
+
 class FieldNotFound(ModelError):
     pass
 

src/couch_potato/model.py

@@ -15,7 +15,7 @@ def format_keys(self):
         return [i[1] for i in Formatter().parse(self.format_str) if i[1] is not None]
 
 
-class Model(BaseModel, metaclass=ModelMeta):
+class Model(BaseModel, metaclass=ModelMeta, add_to_registry=False):
     __read_only__: bool = False
 
     def __eq__(self, other):