Merge pull request #2525 from sopel-irc/SopelIdentifierMemory-casts

memories: test/fix interactions between `SopelIdentifierMemory` ↔ `dict`

sopel/tools/memories.py

@@ -8,10 +8,31 @@
 
 from collections import defaultdict
 import threading
-from typing import Optional
+from typing import Any, Optional, TYPE_CHECKING, Union
 
 from .identifiers import Identifier, IdentifierFactory
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Mapping
+    from typing import Tuple
+
+    MemoryConstructorInput = Union[Mapping[str, Any], Iterable[Tuple[str, Any]]]
+
+
+class _NO_DEFAULT:
+    """Private class to help with overriding C methods like ``dict.pop()``.
+
+    Some Python standard library features are implemented in pure C, and can
+    have a ``null`` default value for certain parameters that is impossible to
+    emulate at the Python layer. This class is our workaround for that.
+
+    .. warning::
+
+        Plugin authors **SHOULD NOT** use this class. It is not part of Sopel's
+        public API.
+
+    """
+
 
 class SopelMemory(dict):
     """A simple thread-safe ``dict`` implementation.
@@ -20,6 +41,12 @@ class SopelMemory(dict):
     them at the same time from different threads, we use a blocking lock in
     ``__setitem__`` and ``__contains__``.
 
+    .. note::
+
+        Unlike the :class:`dict` on which they are based, ``SopelMemory`` and
+        its derivative types do not accept key-value pairs as keyword arguments
+        at construction time.
+
     .. versionadded:: 3.1
         As ``Willie.WillieMemory``
     .. versionchanged:: 4.0
@@ -57,6 +84,12 @@ def __contains__(self, key):
 class SopelMemoryWithDefault(defaultdict):
     """Same as SopelMemory, but subclasses from collections.defaultdict.
 
+    .. note::
+
+        Unlike the :class:`~collections.defaultdict` on which it is based,
+        ``SopelMemoryWithDefault`` does not accept key-value pairs as keyword
+        arguments at construction time.
+
     .. versionadded:: 4.3
         As ``WillieMemoryWithDefault``
     .. versionchanged:: 6.0
@@ -144,14 +177,47 @@ def __init__(
         *args,
         identifier_factory: IdentifierFactory = Identifier,
     ) -> None:
-        super().__init__(*args)
-        self.make_identifier: IdentifierFactory = identifier_factory
+        if len(args) > 1:
+            raise TypeError(
+                'SopelIdentifierMemory expected at most 1 argument, got {}'
+                .format(len(args))
+            )
+
+        self.make_identifier = identifier_factory
         """A factory to transform keys into identifiers."""
 
+        if len(args) == 1:
+            super().__init__(self._convert_keys(args[0]))
+        else:
+            super().__init__()
+
     def _make_key(self, key: Optional[str]) -> Optional[Identifier]:
-        if key is not None:
-            return self.make_identifier(key)
-        return None
+        if key is None:
+            return None
+        return self.make_identifier(key)
+
+    def _convert_keys(
+        self,
+        data: MemoryConstructorInput,
+    ) -> Iterable[tuple[Identifier, Any]]:
+        """Ensure input keys are converted to ``Identifer``.
+
+        :param data: the data passed to the memory at init or update
+        :return: a generator of key-value pairs with the keys converted
+                 to :class:`~.identifiers.Identifier`
+
+        This private method takes input of a mapping or an iterable of key-value
+        pairs and outputs a generator of key-value pairs ready for use in a new
+        or updated :class:`self` instance. It is designed to work with any of the
+        possible ways initial data can be passed to a :class:`dict`, except that
+        ``kwargs`` must be passed to this method as a dictionary.
+        """
+        # figure out what to generate from
+        if hasattr(data, 'items'):
+            data = data.items()
+
+        # return converted input data
+        return ((self.make_identifier(k), v) for k, v in data)
 
     def __getitem__(self, key: Optional[str]):
         return super().__getitem__(self._make_key(key))
@@ -161,3 +227,58 @@ def __contains__(self, key):
 
     def __setitem__(self, key: Optional[str], value):
         super().__setitem__(self._make_key(key), value)
+
+    def setdefault(self, key: str, default=None):
+        return super().setdefault(self._make_key(key), default)
+
+    def __delitem__(self, key: str):
+        super().__delitem__(self._make_key(key))
+
+    def copy(self):
+        return type(self)(self, identifier_factory=self.make_identifier)
+
+    def get(self, key: str, default=_NO_DEFAULT):
+        if default is _NO_DEFAULT:
+            return super().get(self._make_key(key))
+        return super().get(self._make_key(key), default)
+
+    def pop(self, key: str, default=_NO_DEFAULT):
+        if default is _NO_DEFAULT:
+            return super().pop(self._make_key(key))
+        return super().pop(self._make_key(key), default)
+
+    def update(self, maybe_mapping=tuple()):
+        super().update(self._convert_keys(maybe_mapping))
+
+    def __or__(self, other):
+        if not isinstance(other, dict):
+            return NotImplemented
+
+        # self on the left, so other's keys overwrite
+        new = self.copy()
+        new.update(other)
+        return new
+
+    def __ror__(self, other):
+        if not isinstance(other, dict):
+            return NotImplemented
+
+        # self on the right, so keep only new keys from other
+        new = self.copy()
+        new.update((k, v) for k, v in other.items() if k not in self)
+        return new
+
+    def __ior__(self, other):
+        if not isinstance(other, dict):
+            return NotImplemented
+        self.update(other)
+        return self
+
+    def __eq__(self, other):
+        if not isinstance(other, dict):
+            return NotImplemented
+        return super().__eq__(other)
+
+    def __ne__(self, other):
+        ret = self.__eq__(other)
+        return ret if ret is NotImplemented else not ret