added method for caching the PK of a Model

Author: phenobarbital

datamodel/abstract.py

@@ -153,6 +153,7 @@ class ModelMeta(type):
     __fields__: List
     __field_types__: List
     __aliases__: Dict
+    __primary_keys__: List
     # Class-level cache
     _base_class_cache = {}
 
@@ -162,6 +163,8 @@ def _initialize_fields(attrs, annotations, strict):
         _types_local = {}
         _typing_args = {}
         aliases = {}
+        primary_keys = []  # New list to collect primary key fields
+
         for field, _type in annotations.items():
             if isinstance(_type, InitVar) or _type == InitVar:
                 # Skip InitVar fields;
@@ -202,6 +205,10 @@ def _initialize_fields(attrs, annotations, strict):
             df.name = field
             df.type = _type
 
+            # Check for primary_key in field metadata
+            if df.metadata.get("primary_key", False):
+                primary_keys.append(field)
+
             # Cache reflection info so we DON’T need to call
             # get_origin/get_args repeatedly:
             args = get_args(_type)
@@ -281,7 +288,7 @@ def _initialize_fields(attrs, annotations, strict):
             # Assign the field object to the attrs so dataclass can pick it up
             attrs[field] = df
             cols[field] = df
-        return cols, _types_local, _typing_args, aliases
+        return cols, _types_local, _typing_args, aliases, primary_keys
 
     def __new__(cls, name, bases, attrs, **kwargs):  # noqa
         annotations = attrs.get('__annotations__', {})
@@ -295,12 +302,14 @@ def __new__(cls, name, bases, attrs, **kwargs):  # noqa
             _types = cached['types'].copy()
             _typing_args = cached['_typing_args'].copy()
             aliases = cached['aliases'].copy()
+            primary_keys = cached['primary_keys'].copy()
         else:
             # Compute field from Bases:
             cols = OrderedDict()
             _types = {}
             _typing_args = {}
             aliases = {}
+            primary_keys = []
 
             # Step 1: Collect fields from parent classes
             for base in bases:
@@ -311,9 +320,11 @@ def __new__(cls, name, bases, attrs, **kwargs):  # noqa
                     _typing_args |= base.__typing_args__
                 if hasattr(base, '__aliases__'):
                     aliases |= base.__aliases__
+                if hasattr(base, '__primary_keys__'):
+                    primary_keys += base.__primary_keys__
 
             # Now initialize subclass-specific fields
-            new_cols, new_types, new_typing_args, new_aliases = cls._initialize_fields(
+            new_cols, new_types, new_typing_args, new_aliases, nw_primary_keys = cls._initialize_fields(  # noqa
                 attrs, annotations, strict
             )
 
@@ -322,13 +333,15 @@ def __new__(cls, name, bases, attrs, **kwargs):  # noqa
             _types.update(new_types)
             _typing_args.update(new_typing_args)
             aliases.update(new_aliases)
+            primary_keys.extend(nw_primary_keys)
 
             # Store computed results in cache
             cls._base_class_cache[base_key] = {
                 'cols': cols.copy(),
                 'types': _types.copy(),
                 '_typing_args': _typing_args.copy(),
                 'aliases': aliases.copy(),
+                'primary_keys': primary_keys.copy(),
             }
 
         _columns = cols.keys()
@@ -393,11 +406,33 @@ def __new__(cls, name, bases, attrs, **kwargs):  # noqa
         dc.__initialised__ = False
         dc.__field_types__ = _types
         dc.__aliases__ = aliases
+        # Set the primary_keys on the dataclass
+        dc.__primary_keys__ = primary_keys
         dc.__typing_args__ = _typing_args
         dc.modelName = dc.__name__
 
         # Override __setattr__ method
         setattr(dc, "__setattr__", _dc_method_setattr_)
+
+        # Method to get primary keys without further introspection
+        def get_primary_keys(cls):
+            return cls.__primary_keys__
+
+        # Add the method to the class
+        dc.get_primary_keys = classmethod(get_primary_keys)
+
+        def get_primary_key_fields(cls):
+            """Return a dictionary of primary key fields with their types."""
+            return {name: cls.__columns__[name] for name in cls.__primary_keys__}
+
+        dc.get_primary_key_fields = classmethod(get_primary_key_fields)
+
+        def get_primary_key_values(self):
+            """Get primary key values for this instance as a dictionary."""
+            return {key: getattr(self, key) for key in self.__primary_keys__}
+
+        dc.get_primary_key_values = get_primary_key_values
+
         return dc
 
     def __init__(cls, *args, **kwargs) -> None:

datamodel/models.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import contextlib
-from typing import Any, Dict
+from typing import Any, Dict, get_args
 from enum import Enum, EnumMeta
 # Dataclass
 import inspect
@@ -250,13 +250,22 @@ def __collapse_as_values__(
                     )
             # if it's a list, might contain submodels or scalars
             elif isinstance(value, list):
-                if field.origin is list and field.args:
-                    submodel_class = field.args[0]  # The type inside the list
-                    if issubclass(
-                        submodel_class, ModelMixin
-                    ) and not hasattr(submodel_class, name):
-                        out[name] = json_encoder(value)
-                        continue
+                targs = field.args
+                if len(targs) == 2 and targs[1] is type(None):
+                    # Checking if it's Optional[T]
+                    submodel_class = get_args(targs[0])
+                elif field.origin is list and field.args:
+                    submodel_class = field.args[0]
+                else:
+                    submodel_class = None
+                if field.origin is list and submodel_class and (
+                    issubclass(
+                        submodel_class,
+                        ModelMixin
+                    ) and not hasattr(submodel_class, name)
+                ):
+                    out[name] = json_encoder(value)
+                    continue
                 items_out = []
                 for item in value:
                     if isinstance(item, ModelMixin):

datamodel/version.py

@@ -6,7 +6,7 @@
     'simple library based on python +3.8 to use Dataclass-syntax'
     'for interacting with Data'
 )
-__version__ = '0.10.10'
+__version__ = '0.10.11'
 __copyright__ = 'Copyright (c) 2020-2024 Jesus Lara'
 __author__ = 'Jesus Lara'
 __author_email__ = 'jesuslarag@gmail.com'