from __future__ import annotations from typing_extensions import override """Generic persistent, concurrent dictionary-like facility.""" __copyright__ = """ Copyright (C) 2011,2014 Andreas Kloeckner Copyright (C) 2017 Matt Wala """ __license__ = """ Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ import logging import os import pickle import sqlite3 import sys from collections.abc import Callable, Iterator, Mapping from dataclasses import fields as dc_fields, is_dataclass from enum import Enum from typing import Any, TypeVar, cast from warnings import warn from siphash24 import siphash13 from pytools import Hash # noqa: TC001 # Some places are importing it from here. class RecommendedHashNotFoundWarning(UserWarning): pass try: import attrs except ModuleNotFoundError: _HAS_ATTRS = False else: _HAS_ATTRS = True # Used to speed up the comparison in update_for_numpy_scalar and for mock testing _IS_BIGENDIAN = sys.byteorder == "big" logger = logging.getLogger(__name__) # NOTE: not always available so they get hardcoded here SQLITE_BUSY = getattr(sqlite3, "SQLITE_BUSY", 5) SQLITE_CONSTRAINT_PRIMARYKEY = getattr(sqlite3, "SQLITE_CONSTRAINT_PRIMARYKEY", 1555) __doc__ = """ Persistent Hashing and Persistent Dictionaries ============================================== This module contains functionality that allows hashing with keys that remain valid across interpreter invocations, unlike Python's built-in hashes. This module also provides a disk-backed dictionary that uses persistent hashing. .. autoexception:: NoSuchEntryError .. autoexception:: NoSuchEntryCollisionError .. autoexception:: ReadOnlyEntryError .. autoexception:: CollisionWarning .. autoexception:: NanKeyWarning .. autoclass:: KeyBuilder .. autoclass:: PersistentDict .. autoclass:: WriteOncePersistentDict Internal stuff that is only here because the documentation tool wants it ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ .. class:: K A type variable for the key type of a :class:`PersistentDict`. .. class:: V A type variable for the value type of a :class:`PersistentDict`. """ # {{{ key generation class NanKeyWarning(UserWarning): """Warning raised when a NaN is encountered while hashing a key in :class:`KeyBuilder`.""" class KeyBuilder: """A (stateless) object that computes persistent hashes of objects fed to it. Subclassing this class permits customizing the computation of hash keys. This class follows the same general rules as Python's built-in hashing: - Only immutable objects can be hashed. - If two objects compare equal, they must hash to the same value. - Objects with the same hash may or may not compare equal. In addition, hashes computed with :class:`KeyBuilder` have the following properties: - The hash is persistent across interpreter invocations. - The hash is the same across different Python versions and platforms. - The hash is invariant with respect to :envvar:`PYTHONHASHSEED`. - Hashes are computed using functionality from :mod:`hashlib`. Key builders of this type are used by :class:`PersistentDict`, but other uses are entirely allowable. .. automethod:: __call__ .. automethod:: rec .. staticmethod:: new_hash() Return a new hash instance following the protocol of the ones from :mod:`hashlib`. This will permit switching to different hash algorithms in the future. Subclasses are expected to use this to create new hashes. Not doing so is deprecated and may stop working as early as 2022. .. versionadded:: 2021.2 """ # this exists so that we can (conceivably) switch algorithms at some point # down the road new_hash: Callable[..., Hash] = siphash13 def rec(self, key_hash: Hash, key: Any) -> Hash: """ :arg key_hash: the hash object to be updated with the hash of *key*. :arg key: the (immutable) Python object to be hashed. :returns: the updated *key_hash* .. versionchanged:: 2021.2 Now returns the updated *key_hash*. """ digest = getattr(key, "_pytools_persistent_hash_digest", None) if digest is None and not isinstance(key, type): try: method = key.update_persistent_hash except AttributeError: pass else: inner_key_hash = self.new_hash() method(inner_key_hash, self) digest = inner_key_hash.digest() if digest is None: tp = type(key) tname = tp.__name__ method = None try: method = getattr(self, "update_for_"+tname) except AttributeError: if "numpy" in sys.modules: import numpy as np # Hashing numpy dtypes if ( # Handling numpy >= 1.20, for which # type(np.dtype("float32")) -> "dtype[float32]" tname.startswith("dtype[") # Handling numpy >= 1.25, for which # type(np.dtype("float32")) -> "Float32DType" or tname.endswith("DType") ): if isinstance(key, np.dtype): method = self.update_for_specific_dtype # Hashing numpy scalars elif isinstance(key, np.number | np.bool_): # Non-numpy scalars are handled above in the try block. method = self.update_for_numpy_scalar # numpy arrays are mutable so they are not handled in KeyBuilder if method is None: if issubclass(tp, Enum): method = self.update_for_enum elif is_dataclass(tp): method = self.update_for_dataclass elif _HAS_ATTRS and attrs.has(tp): method = self.update_for_attrs if method is not None: inner_key_hash = self.new_hash() method(inner_key_hash, key) digest = inner_key_hash.digest() if digest is None: raise TypeError( f"unsupported type for persistent hash keying: {type(key)}") if not isinstance(key, type): try: object.__setattr__(key, "_pytools_persistent_hash_digest", digest) except AttributeError: pass except TypeError: pass key_hash.update(digest) return key_hash def __call__(self, key: Any) -> str: """Return the hash of *key*.""" key_hash = self.new_hash() self.rec(key_hash, key) return key_hash.hexdigest() # {{{ updaters # NOTE: None of these should be static or classmethods. While Python itself is # perfectly OK with overriding those with 'normal' methods, type checkers # understandably don't like it. def update_for_type(self, key_hash: Hash, key: type) -> None: key_hash.update( f"{key.__module__}.{key.__qualname__}.{key.__name__}".encode()) update_for_ABCMeta = update_for_type def update_for_int(self, key_hash: Hash, key: int) -> None: sz = 8 # Note: this must match the hash for numpy integers, since # np.int64(1) == 1 while True: try: key_hash.update(key.to_bytes(sz, byteorder="little", signed=True)) return except OverflowError: sz *= 2 def update_for_enum(self, key_hash: Hash, key: Enum) -> None: self.update_for_str(key_hash, str(key)) def update_for_bool(self, key_hash: Hash, key: bool) -> None: key_hash.update(str(key).encode("utf8")) def update_for_float(self, key_hash: Hash, key: float) -> None: import math if math.isnan(key): # Also applies to np.nan, float("nan") warn("Encountered a NaN while hashing. Since NaNs compare unequal " "to themselves, the resulting key can not be retrieved from a " "PersistentDict and will lead to a collision error on retrieval.", NanKeyWarning, stacklevel=1) key_hash.update(key.hex().encode("utf8")) def update_for_complex(self, key_hash: Hash, key: float) -> None: key_hash.update(repr(key).encode("utf-8")) def update_for_str(self, key_hash: Hash, key: str) -> None: key_hash.update(key.encode("utf8")) def update_for_bytes(self, key_hash: Hash, key: bytes) -> None: key_hash.update(key) def update_for_tuple(self, key_hash: Hash, key: tuple[Any, ...]) -> None: for obj_i in key: self.rec(key_hash, obj_i) def update_for_frozenset(self, key_hash: Hash, key: frozenset[Any]) -> None: from pytools import unordered_hash unordered_hash( key_hash, (self.rec(self.new_hash(), key_i).digest() for key_i in key), hash_constructor=self.new_hash) update_for_FrozenOrderedSet = update_for_frozenset def update_for_NoneType(self, key_hash: Hash, key: None) -> None: del key key_hash.update(b"") def update_for_dtype(self, key_hash: Hash, key: Any) -> None: key_hash.update(key.str.encode("utf8")) # Handling numpy >= 1.20, for which # type(np.dtype("float32")) -> "dtype[float32]" # Introducing this method allows subclasses to specially handle all those # dtypes. def update_for_specific_dtype(self, key_hash: Hash, key: Any) -> None: key_hash.update(key.str.encode("utf8")) def update_for_numpy_scalar(self, key_hash: Hash, key: Any) -> None: import numpy as np # tobytes() of np.complex256 and np.float128 are non-deterministic, # so we need to use their repr() instead. if hasattr(np, "complex256") and key.dtype == np.dtype("complex256"): key_hash.update(repr(complex(key)).encode("utf8")) elif hasattr(np, "float128") and key.dtype == np.dtype("float128"): key_hash.update(repr(float(key)).encode("utf8")) else: if _IS_BIGENDIAN: key_hash.update(np.array(key).byteswap().tobytes()) else: key_hash.update(np.array(key).tobytes()) def update_for_dataclass(self, key_hash: Hash, key: Any) -> None: self.rec(key_hash, f"{type(key).__qualname__}.{type(key).__name__}") for fld in dc_fields(key): self.rec(key_hash, fld.name) self.rec(key_hash, getattr(key, fld.name, None)) def update_for_attrs(self, key_hash: Hash, key: Any) -> None: self.rec(key_hash, f"{type(key).__qualname__}.{type(key).__name__}") for fld in attrs.fields(key.__class__): self.rec(key_hash, fld.name) self.rec(key_hash, getattr(key, fld.name, None)) def update_for_frozendict(self, key_hash: Hash, key: Mapping[Any, Any]) -> None: from pytools import unordered_hash unordered_hash( key_hash, (self.rec(self.new_hash(), (k, v)).digest() for k, v in key.items()), hash_constructor=self.new_hash) update_for_immutabledict = update_for_frozendict update_for_constantdict = update_for_frozendict update_for_PMap = update_for_frozendict update_for_Map = update_for_frozendict # {{{ date, time, datetime, timezone def update_for_date(self, key_hash: Hash, key: Any) -> None: # 'date' has no timezone information; it is always naive self.rec(key_hash, key.isoformat()) def update_for_time(self, key_hash: Hash, key: Any) -> None: # 'time' should differentiate between naive and aware import datetime # Convert to datetime object self.rec(key_hash, datetime.datetime.combine(datetime.date.min, key)) self.rec(key_hash, "