diff --git a/src/easyscience/base_classes/collection_base.py b/src/easyscience/base_classes/collection_base.py index 3cc0586..0a4e827 100644 --- a/src/easyscience/base_classes/collection_base.py +++ b/src/easyscience/base_classes/collection_base.py @@ -14,6 +14,7 @@ from typing import Tuple from typing import Union +from easyscience.base_classes.new_base import NewBase from easyscience.global_object.undo_redo import NotarizedDict from ..variable.descriptor_base import DescriptorBase @@ -34,7 +35,7 @@ class CollectionBase(BasedBase, MutableSequence): def __init__( self, name: str, - *args: Union[BasedBase, DescriptorBase], + *args: Union[BasedBase, DescriptorBase, NewBase], interface: Optional[InterfaceFactoryTemplate] = None, unique_name: Optional[str] = None, **kwargs, @@ -64,7 +65,7 @@ def __init__( _kwargs[key] = item kwargs = _kwargs for item in list(kwargs.values()) + _args: - if not issubclass(type(item), (DescriptorBase, BasedBase)): + if not issubclass(type(item), (DescriptorBase, BasedBase, NewBase)): raise AttributeError('A collection can only be formed from easyscience objects.') args = _args _kwargs = {} @@ -90,19 +91,19 @@ def __init__( self.interface = interface self._kwargs._stack_enabled = True - def insert(self, index: int, value: Union[DescriptorBase, BasedBase]) -> None: + def insert(self, index: int, value: Union[DescriptorBase, BasedBase, NewBase]) -> None: """ Insert an object into the collection at an index. :param index: Index for EasyScience object to be inserted. :type index: int :param value: Object to be inserted. - :type value: Union[BasedBase, DescriptorBase] + :type value: Union[BasedBase, DescriptorBase, NewBase] :return: None :rtype: None """ t_ = type(value) - if issubclass(t_, (BasedBase, DescriptorBase)): + if issubclass(t_, (BasedBase, DescriptorBase, NewBase)): update_key = list(self._kwargs.keys()) values = list(self._kwargs.values()) # Update the internal dict @@ -117,7 +118,7 @@ def insert(self, index: int, value: Union[DescriptorBase, BasedBase]) -> None: else: raise AttributeError('Only EasyScience objects can be put into an EasyScience group') - def __getitem__(self, idx: Union[int, slice]) -> Union[DescriptorBase, BasedBase]: + def __getitem__(self, idx: Union[int, slice]) -> Union[DescriptorBase, BasedBase, NewBase]: """ Get an item in the collection based on its index. @@ -151,7 +152,7 @@ def __getitem__(self, idx: Union[int, slice]) -> Union[DescriptorBase, BasedBase keys = list(self._kwargs.keys()) return self._kwargs[keys[idx]] - def __setitem__(self, key: int, value: Union[BasedBase, DescriptorBase]) -> None: + def __setitem__(self, key: int, value: Union[BasedBase, DescriptorBase, NewBase]) -> None: """ Set an item via it's index. @@ -163,7 +164,7 @@ def __setitem__(self, key: int, value: Union[BasedBase, DescriptorBase]) -> None if isinstance(value, Number): # noqa: S3827 item = self.__getitem__(key) item.value = value - elif issubclass(type(value), (BasedBase, DescriptorBase)): + elif issubclass(type(value), (BasedBase, DescriptorBase, NewBase)): update_key = list(self._kwargs.keys()) values = list(self._kwargs.values()) old_item = values[key] @@ -233,7 +234,11 @@ def data(self) -> Tuple: def __repr__(self) -> str: return f'{self.__class__.__name__} `{getattr(self, "name")}` of length {len(self)}' - def sort(self, mapping: Callable[[Union[BasedBase, DescriptorBase]], Any], reverse: bool = False) -> None: + def sort( + self, + mapping: Callable[[Union[BasedBase, DescriptorBase, NewBase]], Any], + reverse: bool = False, + ) -> None: """ Sort the collection according to the given mapping.