Fix and extend testing and code fixes

Author: chadell

diffsync/__init__.py

@@ -14,9 +14,8 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-from collections import defaultdict
 from inspect import isclass
-from typing import Callable, ClassVar, Dict, List, Mapping, MutableMapping, Optional, Text, Tuple, Type, Union
+from typing import Callable, ClassVar, Dict, List, Mapping, Optional, Text, Tuple, Type, Union
 
 from pydantic import BaseModel, PrivateAttr
 import structlog  # type: ignore
@@ -25,7 +24,8 @@
 from .enum import DiffSyncModelFlags, DiffSyncFlags, DiffSyncStatus
 from .exceptions import ObjectAlreadyExists, ObjectStoreWrongType, ObjectNotFound
 from .helpers import DiffSyncDiffer, DiffSyncSyncer
-from .store import BaseStore, LocalStore
+from .store import BaseStore
+from .store.local import LocalStore
 
 
 class DiffSyncModel(BaseModel):
@@ -370,20 +370,12 @@ class DiffSync:
     top_level: ClassVar[List[str]] = []
     """List of top-level modelnames to begin from when diffing or synchronizing."""
 
-    _data: MutableMapping[str, MutableMapping[str, DiffSyncModel]]
-    """Defaultdict storing model instances.
-
-    `self._data[modelname][unique_id] == model_instance`
-    """
-
     def __init__(self, name=None, internal_storage_engine=LocalStore):
         """Generic initialization function.
 
         Subclasses should be careful to call super().__init__() if they override this method.
         """
 
-        self._log = structlog.get_logger().new(diffsync=self)
-        
         if isinstance(internal_storage_engine, BaseStore):
             self.store = internal_storage_engine
             self.store.diffsync = self
@@ -436,10 +428,10 @@ def load(self):
     def dict(self, exclude_defaults: bool = True, **kwargs) -> Mapping:
         """Represent the DiffSync contents as a dict, as if it were a Pydantic model."""
         data: Dict[str, Dict[str, Dict]] = {}
-        for modelname in self._data:
+        for modelname in self.store.get_all_model_names():
             data[modelname] = {}
-            for unique_id, model in self._data[modelname].items():
-                data[modelname][unique_id] = model.dict(exclude_defaults=exclude_defaults, **kwargs)
+            for obj in self.store.get_all(modelname):
+                data[obj.get_type()][obj.get_unique_id()] = obj.dict(exclude_defaults=exclude_defaults, **kwargs)
         return data
 
     def str(self, indent: int = 0) -> str:
@@ -571,6 +563,14 @@ def diff_to(
     # Object Storage Management
     # ------------------------------------------------------------------------------
 
+    def get_all_model_names(self):
+        """Get all model names.
+
+        Returns:
+            List[str]: List of model names
+        """
+        return self.store.get_all_model_names()
+
     def get(
         self, obj: Union[Text, DiffSyncModel, Type[DiffSyncModel]], identifier: Union[Text, Mapping]
     ) -> DiffSyncModel:
@@ -659,7 +659,7 @@ def get_or_instantiate(
         Returns:
             Tuple[DiffSyncModel, bool]: Provides the existing or new object and whether it was created or not.
         """
-        return self.store.get_or_instantiate(modle=model, ids=ids, attrs=attrs)
+        return self.store.get_or_instantiate(model=model, ids=ids, attrs=attrs)
 
     def update_or_instantiate(self, model: Type[DiffSyncModel], ids: Dict, attrs: Dict) -> Tuple[DiffSyncModel, bool]:
         """Attempt to update an existing object with provided ids/attrs or instantiate it with provided identifiers and attrs.
@@ -672,7 +672,18 @@ def update_or_instantiate(self, model: Type[DiffSyncModel], ids: Dict, attrs: Di
         Returns:
             Tuple[DiffSyncModel, bool]: Provides the existing or new object and whether it was created or not.
         """
-        return self.store.update_or_instantiate(modle=model, ids=ids, attrs=attrs)
+        return self.store.update_or_instantiate(model=model, ids=ids, attrs=attrs)
+
+    def count(self, modelname=None):
+        """Count how many objects of one model type exist in the backend store.
+
+        Args:
+            modelname (str): The model name to check the number of elements. If not provided, default to all.
+
+        Returns:
+            Int: Number of elements of the model type
+        """
+        return self.store.count(modelname=modelname)
 
 
 # DiffSyncModel references DiffSync and DiffSync references DiffSyncModel. Break the typing loop:

diffsync/store/__init__.py

@@ -1,44 +1,111 @@
-from collections import defaultdict
-from typing import Callable, ClassVar, Dict, List, Mapping, MutableMapping, Optional, Text, Tuple, Type, Union
+"""BaseStore module."""
+from typing import Dict, List, Mapping, Text, Tuple, Type, Union
+import structlog  # type: ignore
 
-from diffsync.exceptions import ObjectNotFound, ObjectAlreadyExists
+from diffsync.exceptions import ObjectNotFound
+
+# TODO: find proper annotation for "DiffSyncModel"  noqa: F821
 
 
 class BaseStore:
+    """Reference store to be implemented in different backends."""
 
-    def __init__(self, diffsync=None, name=None, *args, **kwargs) -> None:
+    def __init__(self, diffsync=None, name=None, *args, **kwargs) -> None:  # pylint: disable=unused-argument
+        """Init method for BaseStore."""
         self.diffsync = diffsync
         self.name = name if name else self.__class__.__name__
+        self._log = structlog.get_logger().new(diffsync=self)
 
     def __str__(self):
+        """Render store name."""
         return self.name
 
-    def get(self, obj: Union[Text, "DiffSyncModel", Type["DiffSyncModel"]], identifier: Union[Text, Mapping]):
+    def get_all_model_names(self):
+        """Get all the model names stored.
+
+        Return:
+            List[str]: List of all the model names.
+        """
+        raise NotImplementedError
+
+    def get(
+        self, obj: Union[Text, "DiffSyncModel", Type["DiffSyncModel"]], identifier: Union[Text, Mapping]  # noqa: F821
+    ):
+        """Get one object from the data store based on its unique id.
+
+        Args:
+            obj: DiffSyncModel class or instance, or modelname string, that defines the type of the object to retrieve
+            identifier: Unique ID of the object to retrieve, or dict of unique identifier keys/values
+
+        Raises:
+            ValueError: if obj is a str and identifier is a dict (can't convert dict into a uid str without a model class)
+            ObjectNotFound: if the requested object is not present
+        """
         raise NotImplementedError
 
-    def get_all(self, obj: Union[Text, "DiffSyncModel", Type["DiffSyncModel"]]) -> List["DiffSyncModel"]:
+    def get_all(self, obj: Union[Text, "DiffSyncModel", Type["DiffSyncModel"]]) -> List["DiffSyncModel"]:  # noqa: F821
+        """Get all objects of a given type.
+
+        Args:
+            obj: DiffSyncModel class or instance, or modelname string, that defines the type of the objects to retrieve
+
+        Returns:
+            List[DiffSyncModel]: List of Object
+        """
         raise NotImplementedError
 
     def get_by_uids(
-        self, uids: List[Text], obj: Union[Text, "DiffSyncModel", Type["DiffSyncModel"]]
-    ) -> List["DiffSyncModel"]:
+        self, uids: List[Text], obj: Union[Text, "DiffSyncModel", Type["DiffSyncModel"]]  # noqa: F821
+    ) -> List["DiffSyncModel"]:  # noqa: F821
+        """Get multiple objects from the store by their unique IDs/Keys and type.
+
+        Args:
+            uids: List of unique id / key identifying object in the database.
+            obj: DiffSyncModel class or instance, or modelname string, that defines the type of the objects to retrieve
+
+        Raises:
+            ObjectNotFound: if any of the requested UIDs are not found in the store
+        """
         raise NotImplementedError
 
-    def remove(self, obj: "DiffSyncModel", remove_children: bool = False):
+    def remove(self, obj: "DiffSyncModel", remove_children: bool = False):  # noqa: F821
+        """Remove a DiffSyncModel object from the store.
+
+        Args:
+            obj (DiffSyncModel): object to remove
+            remove_children (bool): If True, also recursively remove any children of this object
+
+        Raises:
+            ObjectNotFound: if the object is not present
+        """
         raise NotImplementedError
 
-    def add(self, obj: "DiffSyncModel"):
+    def add(self, obj: "DiffSyncModel"):  # noqa: F821
+        """Add a DiffSyncModel object to the store.
+
+        Args:
+            obj (DiffSyncModel): Object to store
+
+        Raises:
+            ObjectAlreadyExists: if a different object with the same uid is already present.
+        """
         raise NotImplementedError
 
-    def update(self, obj: "DiffSyncModel"):
+    def update(self, obj: "DiffSyncModel"):  # noqa: F821
+        """Update a DiffSyncModel object to the store.
+
+        Args:
+            obj (DiffSyncModel): Object to update
+        """
         raise NotImplementedError
 
     def count(self, modelname):
+        """Returns the number of elements of an specific model name."""
         raise NotImplementedError
 
     def get_or_instantiate(
-        self, model: Type["DiffSyncModel"], ids: Dict, attrs: Dict = None
-    ) -> Tuple["DiffSyncModel", bool]:
+        self, model: Type["DiffSyncModel"], ids: Dict, attrs: Dict = None  # noqa: F821
+    ) -> Tuple["DiffSyncModel", bool]:  # noqa: F821
         """Attempt to get the object with provided identifiers or instantiate it with provided identifiers and attrs.
 
         Args:
@@ -63,8 +130,8 @@ def get_or_instantiate(
         return obj, created
 
     def update_or_instantiate(
-        self, model: Type["DiffSyncModel"], ids: Dict, attrs: Dict
-    ) -> Tuple["DiffSyncModel", bool]:
+        self, model: Type["DiffSyncModel"], ids: Dict, attrs: Dict  # noqa: F821
+    ) -> Tuple["DiffSyncModel", bool]:  # noqa: F821
         """Attempt to update an existing object with provided ids/attrs or instantiate it with provided identifiers and attrs.
 
         Args:
@@ -90,158 +157,3 @@ def update_or_instantiate(
                 setattr(obj, attr, value)
 
         return obj, created
-
-
-class LocalStore(BaseStore):
-
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-
-        self._data = defaultdict(dict)
-
-    def get(self, obj: Union[Text, "DiffSyncModel", Type["DiffSyncModel"]], identifier: Union[Text, Mapping]):
-        """Get one object from the data store based on its unique id.
-
-        Args:
-            obj: DiffSyncModel class or instance, or modelname string, that defines the type of the object to retrieve
-            identifier: Unique ID of the object to retrieve, or dict of unique identifier keys/values
-
-        Raises:
-            ValueError: if obj is a str and identifier is a dict (can't convert dict into a uid str without a model class)
-            ObjectNotFound: if the requested object is not present
-        """
-        if isinstance(obj, str):
-            modelname = obj
-            if not hasattr(self, obj):
-                object_class = None
-            else:
-                object_class = getattr(self, obj)
-        else:
-            object_class = obj
-            modelname = obj.get_type()
-
-        if isinstance(identifier, str):
-            uid = identifier
-        elif object_class:
-            uid = object_class.create_unique_id(**identifier)
-        else:
-            raise ValueError(
-                f"Invalid args: ({obj}, {identifier}): "
-                f"either {obj} should be a class/instance or {identifier} should be a str"
-            )
-
-        if uid not in self._data[modelname]:
-            raise ObjectNotFound(f"{modelname} {uid} not present in {str(self)}")
-        return self._data[modelname][uid]
-
-    def get_all(self, obj: Union[Text, "DiffSyncModel", Type["DiffSyncModel"]]) -> List["DiffSyncModel"]:
-        """Get all objects of a given type.
-
-        Args:
-            obj: DiffSyncModel class or instance, or modelname string, that defines the type of the objects to retrieve
-
-        Returns:
-            List[DiffSyncModel]: List of Object
-        """
-        if isinstance(obj, str):
-            modelname = obj
-        else:
-            modelname = obj.get_type()
-
-        return list(self._data[modelname].values())
-
-    def get_by_uids(
-        self, uids: List[Text], obj: Union[Text, "DiffSyncModel", Type["DiffSyncModel"]]
-    ) -> List["DiffSyncModel"]:
-        """Get multiple objects from the store by their unique IDs/Keys and type.
-
-        Args:
-            uids: List of unique id / key identifying object in the database.
-            obj: DiffSyncModel class or instance, or modelname string, that defines the type of the objects to retrieve
-
-        Raises:
-            ObjectNotFound: if any of the requested UIDs are not found in the store
-        """
-        if isinstance(obj, str):
-            modelname = obj
-        else:
-            modelname = obj.get_type()
-
-        results = []
-        for uid in uids:
-            if uid not in self._data[modelname]:
-                raise ObjectNotFound(f"{modelname} {uid} not present in {str(self)}")
-            results.append(self._data[modelname][uid])
-        return results
-
-    def add(self, obj: "DiffSyncModel"):
-        """Add a DiffSyncModel object to the store.
-
-        Args:
-            obj (DiffSyncModel): Object to store
-
-        Raises:
-            ObjectAlreadyExists: if a different object with the same uid is already present.
-        """
-        modelname = obj.get_type()
-        uid = obj.get_unique_id()
-
-        existing_obj = self._data[modelname].get(uid)
-        if existing_obj:
-            if existing_obj is not obj:
-                raise ObjectAlreadyExists(f"Object {uid} already present", obj)
-            # Return so we don't have to change anything on the existing object and underlying data
-            return
-
-        if not obj.diffsync:
-            obj.diffsync = self.diffsync
-
-        self._data[modelname][uid] = obj
-
-    def update(self, obj: "DiffSyncModel"):
-        modelname = obj.get_type()
-        uid = obj.get_unique_id()
-
-        existing_obj = self._data[modelname].get(uid)
-        if existing_obj is obj:
-            return
-
-        self._data[modelname][uid] = obj
-
-    def remove(self, obj: "DiffSyncModel", remove_children: bool = False):
-        """Remove a DiffSyncModel object from the store.
-
-        Args:
-            obj (DiffSyncModel): object to remove
-            remove_children (bool): If True, also recursively remove any children of this object
-
-        Raises:
-            ObjectNotFound: if the object is not present
-        """
-        modelname = obj.get_type()
-        uid = obj.get_unique_id()
-
-        if uid not in self._data[modelname]:
-            raise ObjectNotFound(f"{modelname} {uid} not present in {str(self)}")
-
-        if obj.diffsync:
-            obj.diffsync = None
-
-        del self._data[modelname][uid]
-
-        if remove_children:
-            for child_type, child_fieldname in obj.get_children_mapping().items():
-                for child_id in getattr(obj, child_fieldname):
-                    try:
-                        child_obj = self.get(child_type, child_id)
-                        self.remove(child_obj, remove_children=remove_children)
-                    except ObjectNotFound:
-                        pass
-                        # Since this is "cleanup" code, log an error and continue, instead of letting the exception raise
-                        # self._log.error(f"Unable to remove child {child_id} of {modelname} {uid} - not found!")
-    
-    def count(self, modelname=None):
-        if not modelname:
-            return sum(len(entries) for entries in self._data.values())
-        else:
-            return len(self._data[modelname])