symbolicai 0.21.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

symai/collect/stats.py

@@ -1,21 +1,24 @@
 import json
 import re
+from collections.abc import Callable
 from json import JSONEncoder
-from typing import Any, Callable, List, Optional, Tuple, Type, Union
+from pathlib import Path
+from typing import Any, Union
 
 import numpy as np
 import torch
 
 from ..ops.primitives import OperatorPrimitives
 from ..symbol import Symbol
+from ..utils import UserMessage
 
-SPECIAL_CONSTANT = '__aggregate_'
-EXCLUDE_LIST     = ['_ipython_canary_method_should_not_exist_', '__custom_documentations__']
+SPECIAL_CONSTANT = "__aggregate_"
+EXCLUDE_LIST = ["_ipython_canary_method_should_not_exist_", "__custom_documentations__"]
 
 
 def _normalize_name(name: str) -> str:
     # Replace any character that is not a letter or a number with an underscore
-    normalized_name = re.sub(r'[^a-zA-Z0-9]', '_', name)
+    normalized_name = re.sub(r"[^a-zA-Z0-9]", "_", name)
     return normalized_name.lower()
 
 
@@ -24,31 +27,35 @@ class AggregatorJSONEncoder(JSONEncoder):
         if isinstance(obj, np.ndarray):
             return obj.tolist()
         # drop active from state
-        elif isinstance(obj, Aggregator):
+        if isinstance(obj, Aggregator):
             state = obj.__dict__.copy()
-            state.pop('_raise_error', None)
-            state.pop('_active', None)
-            state.pop('_finalized', None)
-            state.pop('_map', None)
+            state.pop("_raise_error", None)
+            state.pop("_active", None)
+            state.pop("_finalized", None)
+            state.pop("_map", None)
             # drop everything that starts with SPECIAL_CONSTANT
             for key in list(state.keys()):
-                if  not key.startswith(SPECIAL_CONSTANT) and key != '_value' or \
-                    key == '_value' and obj._value == [] or \
-                    key.replace(SPECIAL_CONSTANT, '') in EXCLUDE_LIST:
+                if (
+                    (not key.startswith(SPECIAL_CONSTANT) and key != "_value")
+                    or (key == "_value" and obj._value == [])
+                    or key.replace(SPECIAL_CONSTANT, "") in EXCLUDE_LIST
+                ):
                     state.pop(key, None)
             return state
         return obj.__dict__
 
 
 class Aggregator(Symbol):
-    def __init__(self,
-                 value: Optional[Union["Aggregator", Symbol]] = None,
-                 path: Optional[str] = None,
-                 active: bool = True,
-                 raise_error: bool = False,
-                 *args, **kwargs):
-        super().__init__(*args,
-                         **kwargs)
+    def __init__(
+        self,
+        value: Union["Aggregator", Symbol] | None = None,
+        path: str | None = None,
+        active: bool = True,
+        raise_error: bool = False,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
         # disable nesy engine to avoid side effects
         self.__disable_nesy_engine__ = True
         if value is not None and isinstance(value, Symbol):
@@ -61,27 +68,38 @@ class Aggregator(Symbol):
             elif not isinstance(self._value, (list, tuple)):
                 self._value = [self._value]
         elif value is not None:
-            raise Exception(f'Aggregator object must be of type Aggregator or Symbol! Got: {type(value)}')
+            UserMessage(
+                f"Aggregator object must be of type Aggregator or Symbol! Got: {type(value)}",
+                raise_with=Exception,
+            )
         else:
             self._value = []
-        self._raise_error   = raise_error
-        self._active    = active
+        self._raise_error = raise_error
+        self._active = active
         self._finalized = False
-        self._map       = None
-        self._path      = path
-
-    def __new__(cls, *args,
-            mixin: Optional[bool] = None,
-            primitives: Optional[List[Type]] = [OperatorPrimitives], # only inherit arithmetic primitives
-            callables: Optional[List[Tuple[str, Callable]]] = None,
-            semantic: bool = False,
-            **kwargs) -> "Symbol":
-        return super().__new__(cls, *args,
-                         mixin=mixin,
-                         primitives=primitives,
-                         callables=callables,
-                         semantic=semantic,
-                         **kwargs)
+        self._map = None
+        self._path = path
+
+    def __new__(
+        cls,
+        *args,
+        mixin: bool | None = None,
+        primitives: list[type] | None = None,  # only inherit arithmetic primitives
+        callables: list[tuple[str, Callable]] | None = None,
+        semantic: bool = False,
+        **kwargs,
+    ) -> "Symbol":
+        if primitives is None:
+            primitives = [OperatorPrimitives]
+        return super().__new__(
+            cls,
+            *args,
+            mixin=mixin,
+            primitives=primitives,
+            callables=callables,
+            semantic=semantic,
+            **kwargs,
+        )
 
     def __getattr__(self, name):
         # replace name special characters and spaces with underscores
@@ -93,12 +111,14 @@ class Aggregator(Symbol):
             self._children.append(aggregator)
             # create a new aggregate aggregator
             # named {SPECIAL_CONSTANT}{name} for automatic aggregation
-            self.__dict__[f'{SPECIAL_CONSTANT}{name}'] = aggregator
+            self.__dict__[f"{SPECIAL_CONSTANT}{name}"] = aggregator
             # add also a property with the same name but without the SPECIAL_CONSTANT prefix as a shortcut
-            self.__dict__[name] = self.__dict__[f'{SPECIAL_CONSTANT}{name}']
+            self.__dict__[name] = self.__dict__[f"{SPECIAL_CONSTANT}{name}"]
             return self.__dict__.get(name)
-        elif not self._active and name not in self.__dict__:
-            raise Exception(f'Aggregator object is frozen! No attribute {name} found!')
+        if not self._active and name not in self.__dict__:
+            UserMessage(
+                f"Aggregator object is frozen! No attribute {name} found!", raise_with=Exception
+            )
         return self.__dict__.get(name)
 
     def __setattr__(self, name, value):
@@ -129,32 +149,36 @@ class Aggregator(Symbol):
     def __setstate__(self, state):
         # replace name special characters and spaces with underscores
         # drop active from state
-        state.pop('_raise_error', None)
-        state.pop('_active', None)
-        state.pop('_finalized', None)
-        state.pop('_map', None)
+        state.pop("_raise_error", None)
+        state.pop("_active", None)
+        state.pop("_finalized", None)
+        state.pop("_map", None)
         return super().__setstate__(state)
 
     @staticmethod
     def _set_values(obj, dictionary, parent, strict: bool = True):
         # recursively reconstruct the object
         for key, value in dictionary.items():
-            if isinstance(value, dict):
+            attr_key = key
+            attr_value = value
+            if isinstance(attr_value, dict):
                 if parent is not None:
-                    obj._path = key
-                value = Aggregator._reconstruct(value, parent=parent, strict=strict)
-            if key.startswith(SPECIAL_CONSTANT):
-                key = key.replace(SPECIAL_CONSTANT, '')
-            if key == '_value':
+                    obj._path = attr_key
+                attr_value = Aggregator._reconstruct(attr_value, parent=parent, strict=strict)
+            if attr_key.startswith(SPECIAL_CONSTANT):
+                attr_key = attr_key.replace(SPECIAL_CONSTANT, "")
+            if attr_key == "_value":
                 try:
-                    value = np.asarray(value, dtype=np.float32)
+                    attr_value = np.asarray(attr_value, dtype=np.float32)
                 except Exception as e:
                     if strict:
-                        raise Exception(f'Could not set value of Aggregator object: {obj.path}! ERROR: {e}') from e
-            obj.__setattr__(key, value)
+                        msg = f"Could not set value of Aggregator object: {obj.path}! ERROR: {e}"
+                        UserMessage(msg)
+                        raise Exception(msg) from e
+            obj.__setattr__(attr_key, attr_value)
 
     @staticmethod
-    def _reconstruct(dictionary, parent = None, strict: bool = True):
+    def _reconstruct(dictionary, parent=None, strict: bool = True):
         obj = Aggregator()
         obj._parent = parent
         if parent is not None:
@@ -163,12 +187,12 @@ class Aggregator(Symbol):
         return obj
 
     def __str__(self) -> str:
-        '''
+        """
         Get the string representation of the Symbol object.
 
         Returns:
             str: The string representation of the Symbol object.
-        '''
+        """
         return str(self.entries)
 
     def _to_symbol(self, other) -> Symbol:
@@ -180,14 +204,13 @@ class Aggregator(Symbol):
 
     @property
     def path(self) -> str:
-        path = ''
-        obj  = self
+        path = ""
+        obj = self
         while obj is not None:
             if obj._path is not None:
-                path = obj._path.replace(SPECIAL_CONSTANT, '') + '.' + path
+                path = obj._path.replace(SPECIAL_CONSTANT, "") + "." + path
             obj = obj._parent
-        path = path[:-1] # remove last dot
-        return path
+        return path[:-1]  # remove last dot
 
     def __or__(self, other: Any) -> Any:
         self.add(other)
@@ -211,8 +234,7 @@ class Aggregator(Symbol):
     @property
     def value(self):
         if self.map is not None:
-            res = np.asarray(self.map(np.asarray(self._value, dtype=np.float32)))
-            return res
+            return np.asarray(self.map(np.asarray(self._value, dtype=np.float32)))
         return np.asarray(self._value, dtype=np.float32)
 
     @property
@@ -226,25 +248,26 @@ class Aggregator(Symbol):
     def _set_map_recursively(self, map):
         self._map = map
         for key, value in self.__dict__.items():
-            if isinstance(value, Aggregator) and (not key.startswith('_') or key.startswith(SPECIAL_CONSTANT)):
+            if isinstance(value, Aggregator) and (
+                not key.startswith("_") or key.startswith(SPECIAL_CONSTANT)
+            ):
                 value.map = map
 
     def shape(self):
         if len(self.entries) > 0:
             return np.asarray(self.entries).shape
-        else:
-            return ()
+        return ()
 
     def serialize(self):
         return json.dumps(self, cls=AggregatorJSONEncoder)
 
     def save(self, path: str):
-        with open(path, 'w') as f:
+        with Path(path).open("w") as f:
             json.dump(self, f, cls=AggregatorJSONEncoder)
 
     @staticmethod
     def load(path: str, strict: bool = True):
-        with open(path, 'r') as f:
+        with Path(path).open() as f:
             json_ = json.load(f)
         return Aggregator._reconstruct(json_, strict=strict)
 
@@ -253,45 +276,68 @@ class Aggregator(Symbol):
 
     def add(self, entries):
         # Add entries to the aggregator
-        if not self.active:
-            if self._finalized:
-                raise Exception('Aggregator object is frozen!')
+        if not self.active and self._finalized:
+            UserMessage("Aggregator object is frozen!", raise_with=Exception)
             return
         try:
-            # Append a new entry to the aggregator
-            assert type(entries) in [tuple, list, np.float32, np.float64, np.ndarray, torch.Tensor, int, float, bool, str] or isinstance(entries, Symbol), 'Entries must be a tuple, list, numpy array, torch tensor, integer, float, boolean, string, or Symbol! Got: {}'.format(type(entries))
-            if type(entries) == torch.Tensor:
-                entries = entries.detach().cpu().numpy().astype(np.float32)
-            elif type(entries) in [tuple, list]:
-                entries = np.asarray(entries, dtype=np.float32)
-            elif type(entries) in [int, float]:
-                entries = entries
-            elif type(entries) == bool:
-                entries = int(entries)
-            elif type(entries) == str:
-                entries = Symbol(entries).embedding.astype(np.float32)
-            elif isinstance(entries, Symbol):
-                # use this to avoid recursion on map setter
-                self.add(entries._value)
-                return
-            elif isinstance(entries, Aggregator):
-                self.add(entries.get())
+            processed_entries = self._prepare_entries(entries)
+            if processed_entries is None:
                 return
-
-            if isinstance(entries, np.ndarray) or isinstance(entries, np.float32):
-                entries = entries.squeeze()
-
-            self.entries.append(entries)
+            processed_entries = self._squeeze_entries(processed_entries)
+            self.entries.append(processed_entries)
         except Exception as e:
+            msg = f"Could not add entries to Aggregator object! Please verify type or original error: {e}"
             if self._raise_error:
-                raise Exception(f'Could not add entries to Aggregator object! Please verify type or original error: {e}') from e
-            else:
-                print('Could not add entries to Aggregator object! Please verify type or original error: {}'.format(e))
+                UserMessage(msg)
+                raise Exception(msg) from e
+            UserMessage(msg)
+
+    def _prepare_entries(self, entries):
+        valid_types = (
+            tuple,
+            list,
+            np.float32,
+            np.float64,
+            np.ndarray,
+            torch.Tensor,
+            int,
+            float,
+            bool,
+            str,
+            Symbol,
+        )
+        assert isinstance(entries, valid_types), (
+            f"Entries must be a tuple, list, numpy array, torch tensor, integer, float, boolean, string, or Symbol! Got: {type(entries)}"
+        )
+        if isinstance(entries, torch.Tensor):
+            return entries.detach().cpu().numpy().astype(np.float32)
+        if isinstance(entries, (tuple, list)):
+            return np.asarray(entries, dtype=np.float32)
+        if isinstance(entries, bool):
+            return int(entries)
+        if isinstance(entries, str):
+            return Symbol(entries).embedding.astype(np.float32)
+        if isinstance(entries, Symbol):
+            # Use this to avoid recursion on map setter
+            self.add(entries._value)
+            return None
+        if isinstance(entries, Aggregator):
+            self.add(entries.get())
+            return None
+        return entries
+
+    def _squeeze_entries(self, entries):
+        if isinstance(entries, (np.ndarray, np.float32)):
+            return entries.squeeze()
+        return entries
 
     def keys(self):
         # Get all key names of items that have the SPECIAL_CONSTANT prefix
-        return [key.replace(SPECIAL_CONSTANT, '') for key in self.__dict__.keys() if not key.startswith('_') and \
-                key.replace(SPECIAL_CONSTANT, '') not in EXCLUDE_LIST]
+        return [
+            key.replace(SPECIAL_CONSTANT, "")
+            for key in self.__dict__
+            if not key.startswith("_") and key.replace(SPECIAL_CONSTANT, "") not in EXCLUDE_LIST
+        ]
 
     @property
     def active(self):
@@ -301,7 +347,7 @@ class Aggregator(Symbol):
     @active.setter
     def active(self, value):
         # Set the active status of the aggregator
-        assert isinstance(value, bool), 'Active status must be a boolean! Got: {}'.format(type(value))
+        assert isinstance(value, bool), f"Active status must be a boolean! Got: {type(value)}"
         self._active = value
 
     @property
@@ -312,25 +358,31 @@ class Aggregator(Symbol):
     @finalized.setter
     def finalized(self, value):
         # Set the finalized status of the aggregator
-        assert isinstance(value, bool), 'Finalized status must be a boolean! Got: {}'.format(type(value))
+        assert isinstance(value, bool), f"Finalized status must be a boolean! Got: {type(value)}"
         self._finalized = value
 
     def finalize(self):
         # Finalizes the dynamic creation of the aggregators and freezes the object to prevent further changes
-        self._active     = False
-        self._finalized  = True
+        self._active = False
+        self._finalized = True
+
         def raise_exception(name, value):
-            if name == 'map':
+            if name == "map":
                 self.__setattr__(name, value)
             else:
-                raise Exception('Aggregator object is frozen!')
+                UserMessage("Aggregator object is frozen!", raise_with=Exception)
+
         self.__setattr__ = raise_exception
+
         def get_attribute(*args, **kwargs):
             return self.__dict__.get(*args, **kwargs)
+
         self.__getattr__ = get_attribute
         # Do the same recursively for all properties of type Aggregator
         for key, value in self.__dict__.items():
-            if isinstance(value, Aggregator) and (not key.startswith('_') or key.startswith(SPECIAL_CONSTANT)):
+            if isinstance(value, Aggregator) and (
+                not key.startswith("_") or key.startswith(SPECIAL_CONSTANT)
+            ):
                 value.finalize()
 
     def get(self, *args, **kwargs):
@@ -342,7 +394,7 @@ class Aggregator(Symbol):
     def clear(self):
         # Clear the entries of the aggregator
         if self._finalized:
-            raise Exception('Aggregator object is frozen!')
+            UserMessage("Aggregator object is frozen!", raise_with=Exception)
         self._value = []
 
     def sum(self, axis=0):