adding the generic and other changes to support initialization

Author: jkanche

src/compressed_lists/base.py

@@ -31,7 +31,7 @@ def __init__(
         unlist_data: Any,
         partitioning: Partitioning,
         element_type: Any = None,
-        element_metadata: dict = None,
+        element_metadata: Optional[dict] = None,
         metadata: Optional[dict] = None,
         validate: bool = True,
     ):
@@ -65,7 +65,7 @@ class for the type of elements.
         if validate:
             _validate_data_and_partitions(self._unlist_data, self._partitioning)
 
-    def _define_output(self, in_place: bool = False) -> "Partitioning":
+    def _define_output(self, in_place: bool = False) -> "CompressedList":
         if in_place is True:
             return self
         else:
@@ -207,7 +207,7 @@ def get_names(self) -> Optional[ut.NamedList]:
         """Get the names of list elements."""
         return self._partitioning.get_names()
 
-    def set_names(self, names: Sequence[str], in_place: bool = False) -> "CompressedList":
+    def set_names(self, names: List[str], in_place: bool = False) -> "CompressedList":
         """Set the names of list elements.
 
         names:
@@ -401,7 +401,7 @@ def __getitem__(self, key: Union[int, str, slice]) -> Any:
         """
         # string keys (names)
         if isinstance(key, str):
-            if key not in self.names:
+            if key not in list(self.get_names()):
                 raise KeyError(f"No element named '{key}'.")
             key = list(self.names).index(key)
 
@@ -422,14 +422,14 @@ def __getitem__(self, key: Union[int, str, slice]) -> Any:
             for i in indices:
                 start, end = self._partitioning.get_partition_range(i)
                 result.append(self.extract_range(start, end))
-
-            # Create a new CompressedList from the result
-            return self.__class__.from_list(
+            
+            current_class_const = type(self)
+            return current_class_const.from_list(
                 result, names=[self.names[i] for i in indices] if self.names[0] is not None else None
             )
 
         else:
-            raise TypeError("Index must be int, str, or slice.")
+            raise TypeError("'key' must be int, str, or slice.")
 
     ##################################
     ######>> abstract methods <<######
@@ -460,8 +460,8 @@ def extract_range(self, start: int, end: int) -> Any:
 
     @classmethod
     def from_list(
-        cls, lst: List[Any], names: Optional[Sequence[str]] = None, metadata: dict = None
-    ) -> "CompressedList[Any]":
+        cls, lst: List[Any], names: Optional[Sequence[str]] = None, metadata: Optional[dict] = None
+    ) -> "CompressedList":
         """Create a CompressedList from a regular list.
 
         This method must be implemented by subclasses to handle
@@ -519,7 +519,7 @@ def unlist(self, use_names: bool = True) -> Any:
         """
         return self._unlist_data
 
-    def relist(self, unlist_data: Any) -> "CompressedList[Any]":
+    def relist(self, unlist_data: Any) -> "CompressedList":
         """Create a new `CompressedList` with the same partitioning but different data.
 
         Args:
@@ -531,15 +531,16 @@ def relist(self, unlist_data: Any) -> "CompressedList[Any]":
         """
         _validate_data_and_partitions(unlist_data, self._partitioning)
 
-        return self.__class__(
+        current_class_const = type(self)
+        return current_class_const(
             unlist_data,
             self._partitioning.copy(),
             element_type=self._element_type,
             element_metadata=self._element_metadata.copy(),
             metadata=self._metadata.copy(),
         )
 
-    def extract_subset(self, indices: Sequence[int]) -> "CompressedList[Any]":
+    def extract_subset(self, indices: Sequence[int]) -> "CompressedList":
         """Extract a subset of elements by indices.
 
         Args:
@@ -555,8 +556,8 @@ def extract_subset(self, indices: Sequence[int]) -> "CompressedList[Any]":
                 raise IndexError(f"Index {i} out of range")
 
         # Extract element lengths and names
-        new_lengths = [self.get_element_lengths()[i] for i in indices]
-        new_names = [self.names[i] for i in indices] if self.names[0] is not None else None
+        new_lengths = ut.subset_sequence(self.get_element_lengths(), indices)
+        new_names = ut.subset_sequence(self.names, indices) if self.names is not None else None
 
         # Create new partitioning
         new_partitioning = Partitioning.from_lengths(new_lengths, new_names)
@@ -573,16 +574,16 @@ def extract_subset(self, indices: Sequence[int]) -> "CompressedList[Any]":
         if isinstance(self._unlist_data, np.ndarray):
             new_data = np.concatenate(new_data)
 
-        # Create new compressed list
-        return self.__class__(
+        current_class_const = type(self)
+        return current_class_const(
             new_data,
             new_partitioning,
             element_type=self._element_type,
             element_metadata={k: v for k, v in self._element_metadata.items() if k in indices},
             metadata=self._metadata.copy(),
         )
 
-    def lapply(self, func: Callable) -> "CompressedList[Any]":
+    def lapply(self, func: Callable) -> "CompressedList":
         """Apply a function to each element.
 
         Args:
@@ -593,4 +594,6 @@ def lapply(self, func: Callable) -> "CompressedList[Any]":
             A new CompressedList with the results.
         """
         result = [func(elem) for elem in self]
-        return self.__class__.from_list(result, self.names, self._metadata)
+        
+        current_class_const = type(self)
+        return current_class_const.from_list(result, self.names, self._metadata)

src/compressed_lists/bool_list.py

@@ -1,10 +1,11 @@
-from typing import List, Union
+from typing import List, Optional, Sequence, Union
 
+import numpy as np
 from biocutils.BooleanList import BooleanList
 
 from .base import CompressedList
 from .partition import Partitioning
-from .split_generic import splitAsCompressedList
+from .split_generic import _generic_register_helper, splitAsCompressedList
 
 __author__ = "Jayaram Kancherla"
 __copyright__ = "Jayaram Kancherla"
@@ -18,8 +19,8 @@ def __init__(
         self,
         unlist_data: BooleanList,
         partitioning: Partitioning,
-        element_metadata: dict = None,
-        metadata: dict = None,
+        element_metadata: Optional[dict] = None,
+        metadata: Optional[dict] = None,
         **kwargs,
     ):
         """Initialize a CompressedBooleanList.
@@ -51,10 +52,47 @@ def __init__(
             unlist_data, partitioning, element_type=BooleanList, element_metadata=element_metadata, metadata=metadata
         )
 
+    @classmethod
+    def from_partitioned_data(
+        cls, partitioned_data: List[List], partitioning: Partitioning, metadata: Optional[dict] = None
+    ) -> "CompressedBooleanList":
+        """Create `CompressedBooleanList` from already-partitioned data.
+
+        Args:
+            partitioned_data:
+                List of `BooleanList`'s, each containing booleans for one partition.
+
+            partitioning:
+                Partitioning object defining the boundaries.
+
+            metadata:
+                Optional metadata.
+
+        Returns:
+            A new `CompressedBooleanList`.
+        """
+        flat_data = []
+        for partition in partitioned_data:
+            flat_data.extend(partition)
+
+        unlist_data = BooleanList(flat_data)
+
+        return cls(unlist_data, partitioning, metadata=metadata)
+
 
 @splitAsCompressedList.register
 def _(
-    data: Union[List[List[bool]], List[BooleanList]], names: List[str] = None, metadata: dict = None
+    data: BooleanList,
+    groups_or_partitions: Union[list, Partitioning],
+    names: Optional[Sequence[str]] = None,
+    metadata: Optional[dict] = None,
 ) -> CompressedBooleanList:
-    """Handle lists of boolean."""
-    return CompressedBooleanList.from_list(data, names, metadata)
+    """Handle lists of booleans."""
+
+    partitioned_data, groups_or_partitions = _generic_register_helper(
+        data=data, groups_or_partitions=groups_or_partitions, names=names
+    )
+
+    return CompressedBooleanList.from_partitioned_data(
+        partitioned_data=partitioned_data, partitioning=groups_or_partitions, metadata=metadata
+    )

src/compressed_lists/float_list.py

@@ -1,10 +1,11 @@
-from typing import List, Union
+from typing import List, Optional, Sequence, Union
 
+import numpy as np
 from biocutils.FloatList import FloatList
 
 from .base import CompressedList
 from .partition import Partitioning
-from .split_generic import splitAsCompressedList
+from .split_generic import _generic_register_helper, splitAsCompressedList
 
 __author__ = "Jayaram Kancherla"
 __copyright__ = "Jayaram Kancherla"
@@ -18,8 +19,8 @@ def __init__(
         self,
         unlist_data: FloatList,
         partitioning: Partitioning,
-        element_metadata: dict = None,
-        metadata: dict = None,
+        element_metadata: Optional[dict] = None,
+        metadata: Optional[dict] = None,
         **kwargs,
     ):
         """Initialize a CompressedFloatList.
@@ -51,10 +52,47 @@ def __init__(
             unlist_data, partitioning, element_type=FloatList, element_metadata=element_metadata, metadata=metadata
         )
 
+    @classmethod
+    def from_partitioned_data(
+        cls, partitioned_data: List[List], partitioning: Partitioning, metadata: Optional[dict] = None
+    ) -> "CompressedFloatList":
+        """Create `CompressedFloatList` from already-partitioned data.
+
+        Args:
+            partitioned_data:
+                List of lists, each containing floats for one partition.
+
+            partitioning:
+                Partitioning object defining the boundaries.
+
+            metadata:
+                Optional metadata.
+
+        Returns:
+            A new `CompressedFloatList`.
+        """
+        flat_data = []
+        for partition in partitioned_data:
+            flat_data.extend(partition)
+
+        unlist_data = FloatList(flat_data)
+
+        return cls(unlist_data, partitioning, metadata=metadata)
+
 
 @splitAsCompressedList.register
 def _(
-    data: Union[List[List[float]], List[FloatList]], names: List[str] = None, metadata: dict = None
+    data: FloatList,
+    groups_or_partitions: Union[list, Partitioning],
+    names: Optional[Sequence[str]] = None,
+    metadata: Optional[dict] = None,
 ) -> CompressedFloatList:
     """Handle lists of floats."""
-    return CompressedFloatList.from_list(data, names, metadata)
+
+    partitioned_data, groups_or_partitions = _generic_register_helper(
+        data=data, groups_or_partitions=groups_or_partitions, names=names
+    )
+
+    return CompressedFloatList.from_partitioned_data(
+        partitioned_data=partitioned_data, partitioning=groups_or_partitions, metadata=metadata
+    )

src/compressed_lists/integer_list.py

@@ -1,10 +1,11 @@
-from typing import List, Union
+from typing import Optional, Sequence, Union
 
+import numpy as np
 from biocutils.IntegerList import IntegerList
 
 from .base import CompressedList
 from .partition import Partitioning
-from .split_generic import splitAsCompressedList
+from .split_generic import splitAsCompressedList, _generic_register_helper
 
 __author__ = "Jayaram Kancherla"
 __copyright__ = "Jayaram Kancherla"
@@ -18,8 +19,8 @@ def __init__(
         self,
         unlist_data: IntegerList,
         partitioning: Partitioning,
-        element_metadata: dict = None,
-        metadata: dict = None,
+        element_metadata: Optional[dict] = None,
+        metadata: Optional[dict] = None,
         **kwargs,
     ):
         """Initialize a CompressedIntegerList.
@@ -51,10 +52,48 @@ def __init__(
             unlist_data, partitioning, element_type=IntegerList, element_metadata=element_metadata, metadata=metadata
         )
 
+    @classmethod
+    def from_partitioned_data(
+        cls, partitioned_data: Sequence[IntegerList], partitioning: Partitioning, metadata: Optional[dict] = None
+    ) -> "CompressedIntegerList":
+        """Create `CompressedIntegerList` from already-partitioned data.
+
+        Args:
+            partitioned_data:
+                List of `IntegerList`'s, each containing integers for one partition.
+
+            partitioning:
+                Partitioning object defining the boundaries.
+
+            metadata:
+                Optional metadata.
+
+        Returns:
+            A new `CompressedIntegerList`.
+        """
+
+        flat_data = []
+        for partition in partitioned_data:
+            flat_data.extend(partition)
+
+        unlist_data = IntegerList(flat_data)
+
+        return cls(unlist_data, partitioning, metadata=metadata)
+
 
 @splitAsCompressedList.register
 def _(
-    data: Union[List[List[int]], List[IntegerList]], names: List[str] = None, metadata: dict = None
+    data: IntegerList,
+    groups_or_partitions: Union[list, Partitioning],
+    names: Optional[Sequence[str]] = None,
+    metadata: Optional[dict] = None,
 ) -> CompressedIntegerList:
     """Handle lists of integers."""
-    return CompressedIntegerList.from_list(data, names, metadata)
+    
+    partitioned_data, groups_or_partitions = _generic_register_helper(
+        data=data, groups_or_partitions=groups_or_partitions, names=names
+    )
+    
+    return CompressedIntegerList.from_partitioned_data(
+        partitioned_data=partitioned_data, partitioning=groups_or_partitions, metadata=metadata
+    )