datamaestro 1.0.5__py3-none-any.whl → 1.0.6__py3-none-any.whl

datamaestro/record.py

@@ -1,4 +1,5 @@
-from typing import ClassVar, Type, TypeVar, Dict, List, Union, Optional
+import logging
+from typing import ClassVar, Type, TypeVar, Dict, List, Union, Optional, FrozenSet
 
 
 class Item:
@@ -6,6 +7,7 @@ class Item:
 
     @classmethod
     def __get_base__(cls: Type) -> Type:
+        """Get the most generic superclass for this type of item"""
         if base := getattr(cls, "__base__cache__", None):
             return base
 
@@ -22,33 +24,36 @@ Items = Dict[Type[T], T]
 
 
 class Record:
-    """Associate types with entries"""
+    """Associate types with entries
+
+    A record is a composition of items; each item base class is unique.
+    """
 
     items: Items
 
-    def __init__(self, *items: Union[Items, T], no_check=False):
+    def __init__(self, *items: Union[Items, T], override=False):
         self.items = {}
 
         if len(items) == 1 and isinstance(items[0], dict):
+            # Just copy the dictionary
             self.items = items[0]
         else:
-            for item in items:
-                self._add(item, update_only=True)
-
-        # Check if the record is constructured
-        if not no_check:
-            self.validate()
+            for entry in items:
+                # Returns a new record if the item exists
+                base = entry.__get_base__()
+                if not override and base in self.items:
+                    raise RuntimeError(
+                        f"The item type {base} ({entry.__class__})"
+                        " is already in the record"
+                    )
+                self.items[base] = entry
 
-    def __new__(cls, *items: Union[Items, T], no_check=False):
-        # Without this, impossible to pickle objects
-        if cls.__trueclass__ is not None:
-            record = object.__new__(cls.__trueclass__)
-            record.__init__(*items, no_check=True)
-            if not no_check:
-                record.validate(cls=cls)
-            return record
+        self.validate()
 
-        return object.__new__(cls)
+    @classmethod
+    def from_record(cls, record: "Record", *items: T, override=True):
+        """Build from another record"""
+        return cls({**record.items, **{item.__get_base__(): item for item in items}})
 
     def __str__(self):
         return (
@@ -57,6 +62,21 @@ class Record:
             + "}"
         )
 
+    def __reduce__(self):
+        cls = self.__class__
+        if cls.__trueclass__ is None:
+            return (cls.__new__, (cls.__trueclass__ or cls,), {"items": self.items})
+
+        return (
+            cls.__new__,
+            (cls.__trueclass__ or cls,),
+            {"items": self.items, "itemtypes": self.itemtypes},
+        )
+
+    def __setstate__(self, state):
+        self.items = state["items"]
+        self.itemtypes = None
+
     def validate(self, cls: Type["Record"] = None):
         """Validate the record"""
         cls = cls if cls is not None else self.__class__
@@ -77,17 +97,19 @@ class Record:
                     for itemtype in cls.itemtypes
                 )
             ]
-            raise RuntimeError(
+            raise KeyError(
                 f"The record {cls} contains unregistered items: {unregistered}"
             )
 
     def get(self, key: Type[T]) -> Optional[T]:
+        """Get a given item or None if it does not exist"""
         try:
             return self[key]
         except KeyError:
             return None
 
     def has(self, key: Type[T]) -> bool:
+        """Returns True if the record has the given item type"""
         return key.__get_base__() in self.items
 
     def __getitem__(self, key: Type[T]) -> T:
@@ -100,35 +122,27 @@ class Record:
             raise KeyError(f"No entry with type {key}")
         return entry
 
-    def add(self, *entries: T, update_only=False, no_check=False) -> "Record":
-        """Update the record with these new items, and returns a new record if
-        any item already exists"""
-        return self._add(*entries, update_only=update_only)
-
-    def _add(self, *entries: T, update_only=False, no_check=False) -> "Record":
-        """Internal method for updating records"""
-        for entry in entries:
-            # Returns a new record if the item exists
-            base = entry.__get_base__()
-            if base in self.items:
-                if update_only:
-                    raise RuntimeError(
-                        f"The item type {base} ({entry.__class__})"
-                        " is already in the record"
-                    )
-                return self.__class__({**self.items, base: entry}, no_check=no_check)
+    def is_pickled(self):
+        return self.itemtypes is None
 
-            # No, just update
-            self.items[base] = entry
-        return self
+    def update(self, *items: T) -> "Record":
+        """Update some items"""
+        # Create our new dictionary
+        item_dict = {**self.items}
+        for item in items:
+            item_dict[item.__get_base__()] = item
+
+        return self.__class__(item_dict)
 
     # --- Class methods and variables
 
-    itemtypes: ClassVar[List[Type[T]]] = []
-    """For specific records, this is the list of types"""
+    itemtypes: ClassVar[Optional[FrozenSet[Type[T]]]] = []
+    """For specific records, this is the list of types. The value is null when
+    no validation is used (e.g. pickled records created on the fly)"""
 
     __trueclass__: ClassVar[Optional[Type["Record"]]] = None
-    """True when the class is defined in a module"""
+    """The last class in the type hierarchy corresponding to an actual type,
+    i.e. not created on the fly (only defined when the record is pickled)"""
 
     @classmethod
     def has_type(cls, itemtype: Type[T]):
@@ -138,31 +152,52 @@ class Record:
     def _subclass(cls, *itemtypes: Type[T]):
         cls_itemtypes = [x for x in getattr(cls, "itemtypes", [])]
         mapping = {
-            ix: itemtype.__get_base__() for ix, itemtype in enumerate(cls_itemtypes)
+            itemtype.__get_base__(): ix for ix, itemtype in enumerate(cls_itemtypes)
         }
 
         for itemtype in itemtypes:
-            if ix := mapping.get(itemtype.__get_base__(), None):
+            if (ix := mapping.get(itemtype.__get_base__(), -1)) >= 0:
                 cls_itemtypes[ix] = itemtype
             else:
                 cls_itemtypes.append(itemtype)
-        return cls_itemtypes
+
+        return frozenset(cls_itemtypes)
 
     @classmethod
     def from_types(cls, name: str, *itemtypes: Type[T], module: str = None):
+        """Construct a new sub-record type
+
+        :param name: The name of the subrecord
+        :param module: The module name, defaults to None
+        :return: A new Record type
+        """
         extra_dict = {}
         if module:
             extra_dict["__module__"] = module
+
         return type(
             name,
             (cls,),
             {
                 **extra_dict,
+                "itemtypes": frozenset(cls._subclass(*itemtypes)),
                 "__trueclass__": cls.__trueclass__ or cls,
-                "itemtypes": cls._subclass(*itemtypes),
             },
         )
 
+    __RECORD_TYPES_CACHE__: Dict[frozenset, Type["Record"]] = {}
+
+    @staticmethod
+    def fromitemtypes(itemtypes: FrozenSet[T]):
+        if recordtype := Record.__RECORD_TYPES_CACHE__.get(itemtypes, None):
+            return recordtype
+
+        recordtype = Record.from_types(
+            "_".join(itemtype.__name__ for itemtype in itemtypes), *itemtypes
+        )
+        Record.__RECORD_TYPES_CACHE__[itemtypes] = recordtype
+        return recordtype
+
 
 def recordtypes(*types: List[Type[T]]):
     """Adds types for a new record class"""
@@ -176,27 +211,102 @@ def recordtypes(*types: List[Type[T]]):
     return decorate
 
 
-class RecordTypesCache:
-    """Class to use when new record types need to be created on the fly by
-    adding new items"""
-
+class RecordTypesCacheBase:
     def __init__(self, name: str, *itemtypes: Type[T], module: str = None):
         self._module = module
         self._name = name
         self._itemtypes = itemtypes
-        self._cache: Dict[Type[Record], Type[Record]] = {}
-
-    def __getitem__(self, record_type: Type[Record]):
-        if updated_type := self._cache.get(record_type, None):
-            return updated_type
 
+    def _compute(self, record_type: Type[Record]):
         updated_type = record_type.from_types(
             f"{self._name}_{record_type.__name__}",
             *self._itemtypes,
             module=self._module,
         )
-        self._cache[record_type] = updated_type
         return updated_type
 
-    def update(self, record: Record, *items: Item):
-        return self[record.__class__](*record.items.values(), *items)
+
+class RecordTypesCache(RecordTypesCacheBase):
+    """Class to use when new record types need to be created on the fly by
+    adding new items"""
+
+    def __init__(self, name: str, *itemtypes: Type[T], module: str = None):
+        """Creates a new cache
+
+        :param name: Base name for new record types
+        :param module: The module name for new types, defaults to None
+        """
+        super().__init__(name, *itemtypes, module=module)
+        self._cache: Dict[Type[Record], Type[Record]] = {}
+        self._warning = False
+
+    def __call__(self, record_type: Type[Record]):
+        if (updated_type := self._cache.get(record_type, None)) is None:
+            self._cache[record_type] = updated_type = self._compute(record_type)
+        return updated_type
+
+    def update(self, record: Record, *items: Item, cls=None):
+        """Update the record with the given items
+
+        :param record: The record to which we add items
+        :param cls: The class of the record, useful if the record has been
+            pickled, defaults to None
+        :return: A new record with the extra items
+        """
+        if cls is None:
+            cls = record.__class__
+            if record.is_pickled() and not self._warning:
+                logging.warning(
+                    "Updating unpickled records is not recommended"
+                    " (speed issues): use the pickle record class as the cls input"
+                )
+                itemtypes = frozenset(type(item) for item in record.items.values())
+                cls = Record.fromitemtypes(itemtypes)
+        else:
+            assert (
+                record.is_pickled()
+            ), "cls can be used only when the record as been pickled"
+
+        return self(cls)(*record.items.values(), *items, override=True)
+
+
+class SingleRecordTypeCache(RecordTypesCacheBase):
+    """Class to use when new record types need to be created on the fly by
+    adding new items
+
+    This class supposes that the input record type is always the same (no check
+    is done to ensure this)"""
+
+    def __init__(self, name: str, *itemtypes: Type[T], module: str = None):
+        """Creates a new cache
+
+        :param name: Base name for new record types
+        :param module: The module name for new types, defaults to None
+        """
+        super().__init__(name, *itemtypes, module=module)
+        self._cache: Optional[Type[Record]] = None
+
+    def __call__(self, record_type: Type[Record]):
+        if self._cache is None:
+            self._cache = self._compute(record_type)
+        return self._cache
+
+    def update(self, record: Record, *items: Item, cls=None):
+        """Update the record with the given items
+
+        :param record: The record to which we add items
+        :param cls: The class of the record, useful if the record has been
+            pickled, defaults to None
+        :return: A new record with the extra items
+        """
+        if self._cache is None:
+            if cls is None:
+                cls = record.__class__
+                itemtypes = frozenset(type(item) for item in record.items.values())
+                cls = Record.fromitemtypes(itemtypes)
+            else:
+                assert (
+                    record.is_pickled()
+                ), "cls can be used only when the record as been pickled"
+
+        return self(cls)(*record.items.values(), *items, override=True)

datamaestro/test/test_record.py

@@ -1,4 +1,11 @@
-from datamaestro.record import Record, Item, RecordTypesCache, recordtypes
+import pickle
+from datamaestro.record import (
+    Record,
+    Item,
+    RecordTypesCache,
+    recordtypes,
+    SingleRecordTypeCache,
+)
 from attrs import define
 import pytest
 
@@ -18,13 +25,24 @@ class BItem(Item):
     b: int
 
 
+@define
+class B1Item(BItem):
+    b1: int
+
+
 @define
 class CItem(Item):
     c: int
 
 
-class MyRecord(Record):
-    itemtypes = [A1Item, BItem]
+@recordtypes(A1Item)
+class BaseRecord(Record):
+    ...
+
+
+@recordtypes(BItem)
+class MyRecord(BaseRecord):
+    ...
 
 
 @recordtypes(CItem)
@@ -43,6 +61,7 @@ def test_record_simple():
 
 def test_record_missing_init():
     with pytest.raises(KeyError):
+        # A1Item is missing
         MyRecord(AItem(1), BItem(2))
 
     with pytest.raises(KeyError):
@@ -54,33 +73,79 @@ def test_record_update():
     b = BItem(4)
     r = MyRecord(a, b)
 
-    r2 = r.add(BItem(3))
+    r2 = r.update(BItem(3))
     assert r is not r2
     assert r2[BItem] is not b
 
+    r3 = MyRecord2.from_record(r, CItem(2), BItem(5))
+    assert r[BItem].b == 4
+    assert r3[BItem].b == 5
+
 
 def test_record_decorator():
     MyRecord2(A1Item(1, 2), BItem(2), CItem(3))
 
 
-def test_record_newtype():
-    MyRecord2 = MyRecord.from_types("MyRecord2", CItem)
-    r = MyRecord2(A1Item(1, 2), BItem(2), CItem(3))
-
-    # For a dynamic class, we should have the same MyRecord type
-    assert r.__class__ is MyRecord
+def test_record_type_update():
+    itemtypes = MyRecord2.from_types("Test", B1Item).itemtypes
+    assert itemtypes == frozenset((A1Item, B1Item, CItem))
 
 
 def test_record_onthefly():
     cache = RecordTypesCache("OnTheFly", CItem)
 
-    MyRecord2 = cache[MyRecord]
-    r2 = MyRecord2(A1Item(1, 2), BItem(2), CItem(3))
-    assert r2.__class__ is MyRecord
+    MyRecord2 = cache(MyRecord)
+    MyRecord2(A1Item(1, 2), BItem(2), CItem(3))
 
-    assert cache[MyRecord] is MyRecord2
+    assert cache(MyRecord) is MyRecord2
 
     r = MyRecord(A1Item(1, 2), BItem(2))
-    assert cache[r.__class__] is MyRecord2
+    assert cache(r.__class__) is MyRecord2
+
+    r = cache.update(r, CItem(3))
+
+    # Same record type
+    cache2 = RecordTypesCache("OnTheFly", CItem)
+
+    cache2.update(r, CItem(4))
+
+
+def test_record_pickled():
+    # First,
+    MyRecord2 = BaseRecord.from_types("MyRecordBis", BItem)
+    r = MyRecord2(A1Item(1, 2), BItem(2))
+    r = pickle.loads(pickle.dumps(r))
+
+    assert isinstance(r, BaseRecord) and not isinstance(r, MyRecord2)
+    cache = RecordTypesCache("OnTheFly", CItem)
+
+    assert r.is_pickled()
+
+    r2 = cache.update(r, CItem(4))
+    assert not r2.is_pickled()
+
+    # Test with cls update
+    with pytest.raises(KeyError):
+        cache.update(r, CItem(4), cls=BaseRecord)
+
+    # This is OK
+    cache.update(r, CItem(4), cls=MyRecord)
+
+    # --- Test when we update a pickled record with an of a sub-class
+    cache = RecordTypesCache("OnTheFly", B1Item)
+    r2 = cache.update(r, B1Item(1, 2))
+
+
+def test_record_pickled_single():
+    MyRecord2 = BaseRecord.from_types("MyRecordBis", BItem)
+    r = MyRecord2(A1Item(1, 2), BItem(2))
+    r = pickle.loads(pickle.dumps(r))
+
+    cache = SingleRecordTypeCache("OnTheFly", CItem)
+
+    updated = cache.update(r, CItem(4))
+
+    assert updated.itemtypes == frozenset((A1Item, BItem, CItem))
 
-    cache.update(r, CItem(3))
+    # Even with the wrong record, no change now
+    assert cache(BaseRecord).itemtypes == frozenset((A1Item, BItem, CItem))

datamaestro/version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.0.5'
-__version_tuple__ = version_tuple = (1, 0, 5)
+__version__ = version = '1.0.6'
+__version_tuple__ = version_tuple = (1, 0, 6)

{datamaestro-1.0.5.dist-info → datamaestro-1.0.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datamaestro
-Version: 1.0.5
+Version: 1.0.6
 Summary: "Dataset management command line and API"
 Home-page: https://github.com/experimaestro/datamaestro
 Author: Benjamin Piwowarski

{datamaestro-1.0.5.dist-info → datamaestro-1.0.6.dist-info}/RECORD

@@ -2,13 +2,13 @@ datamaestro/__init__.py,sha256=9M5hA6FVngduJBcjInvJWQM8n0cqapXAFPzfRLHR74c,237
 datamaestro/__main__.py,sha256=tJTf1sTWKRIatvBcHlWDIZRZodAZ2B2zkD01pD89MYk,9024
 datamaestro/context.py,sha256=imDAs9v5yR_O871cdM4e5Q0KBSHmT-j5qc83sz6npWI,13210
 datamaestro/definitions.py,sha256=ORlD3kxLvmw1EScCepzGQOHjsTiC3yvY-z_7LxKDp08,15415
-datamaestro/record.py,sha256=xFFyXRbV_cm52o3KUDFN0G7TseRgqO7IpFfIAJlmCh0,6386
+datamaestro/record.py,sha256=gGOJAKHKc3SN3zctVZ6Bq21cuLnWymc2cE9ukeyrZzs,10475
 datamaestro/registry.py,sha256=M7QJkcWJP_cxAoqIioLQ01ou2Zg9RqGQvW0XGVspYFE,1421
 datamaestro/search.py,sha256=PMceNp5hcp0dlzs4cLb6LJT7XHrdXo58oO7oTucawbE,2887
 datamaestro/settings.py,sha256=HYSElTUYZ6DZocBb9o3ifm6WW9knRO64XJUwxGIpvwQ,1304
 datamaestro/sphinx.py,sha256=bp7x_2BFoTSwTqcVZDM8R8cWa7G2pz0Zb8GS054lLYM,6996
 datamaestro/utils.py,sha256=Y3_aqeOHW8vuifwggGWJfgONyDG1FLX7ONAnX85jENI,6511
-datamaestro/version.py,sha256=ClFIlbf5O23dVWC6oD6t3m2n6nB22IOt3V9t-YX_mG0,411
+datamaestro/version.py,sha256=T17ZEPR5Omt5RLA0TuZWTufRi58PzGIPek75nFcuMQY,411
 datamaestro/annotations/__init__.py,sha256=jLprrxSBa5QIqc--vqycEcxU4CR9WjVNRaqR5lH0EuE,39
 datamaestro/annotations/agreement.py,sha256=IPHjXX8ld5blvSDNXGt8RfuHDgVAITN52gJZh93AY4g,723
 datamaestro/commands/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -37,10 +37,10 @@ datamaestro/test/checks.py,sha256=1eTkz4YJhAPOcnQSsz4vPnvzwwfrEnpn6H_s1ADISpo,17
 datamaestro/test/conftest.py,sha256=it4S5Qq1CA_U8qM0pr4m7v-1dhLj5Y49WjVg5Ee3mpM,767
 datamaestro/test/test_annotations.py,sha256=kRPUmS_UAN6JSSVPUwV4OM_LEuEUHF1OcLSiYXjsKjw,246
 datamaestro/test/test_download_handlers.py,sha256=Qqm-fML1KVp6dPwAUcH6xzi_dpQIshvROzviSYCUzc0,603
-datamaestro/test/test_record.py,sha256=78wgQoVyXiej0MmaDwHPTlqXh5bqlWU6AHTTY89qLPU,1571
-datamaestro-1.0.5.dist-info/LICENSE,sha256=WJ7YI-moTFb-uVrFjnzzhGJrnL9P2iqQe8NuED3hutI,35141
-datamaestro-1.0.5.dist-info/METADATA,sha256=7dr7-Q8ld_KY5IUkv3dwKamBmG0vqns1u6wbW9wXEPA,8999
-datamaestro-1.0.5.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-datamaestro-1.0.5.dist-info/entry_points.txt,sha256=8qMhwSRvFG2iBqtJYVD22Zd4s4c3YkODtcp0Ajw1knw,133
-datamaestro-1.0.5.dist-info/top_level.txt,sha256=XSznaMNAA8jELV7-TOqaAgDsjLzUf9G9MxL7C4helT0,12
-datamaestro-1.0.5.dist-info/RECORD,,
+datamaestro/test/test_record.py,sha256=zo8faY8xqRRE916mXIGaq-WXhhvQrNFmDJTzB3CCFEk,3050
+datamaestro-1.0.6.dist-info/LICENSE,sha256=WJ7YI-moTFb-uVrFjnzzhGJrnL9P2iqQe8NuED3hutI,35141
+datamaestro-1.0.6.dist-info/METADATA,sha256=U_fiwzaUSjEV5s4DOp3dRn9MzXesAkQ2o_ZpXCZIW8Y,8999
+datamaestro-1.0.6.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+datamaestro-1.0.6.dist-info/entry_points.txt,sha256=8qMhwSRvFG2iBqtJYVD22Zd4s4c3YkODtcp0Ajw1knw,133
+datamaestro-1.0.6.dist-info/top_level.txt,sha256=XSznaMNAA8jELV7-TOqaAgDsjLzUf9G9MxL7C4helT0,12
+datamaestro-1.0.6.dist-info/RECORD,,