feat: add SubsetFilterProxyModel

prettyqt/core/abstractitemmodel.py

@@ -22,11 +22,40 @@
 )
 
 
+class Subset:
+    def __init__(self, model):
+        self.model = model
+
+    def __getitem__(self, index):
+        parent = self.model.parent()
+        if parent is None:
+            raise ValueError("needs parent!")
+
+        from prettyqt import custom_models
+
+        match index:
+            case (arg_1, arg_2):
+                proxy = custom_models.SubsetFilterProxyModel(
+                    row_filter=arg_1, column_filter=arg_2, parent=parent
+                )
+            case _:
+                proxy = custom_models.SubsetFilterProxyModel(
+                    row_filter=index, column_filter=None, parent=parent
+                )
+                # raise ValueError(index)
+        proxy.setSourceModel(self.model)
+        return proxy
+
+
 class AbstractItemModelMixin(core.ObjectMixin):
     DEFAULT_FLAGS = (
         constants.DRAG_ENABLED | constants.IS_ENABLED | constants.IS_SELECTABLE
     )
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.subset = Subset(self)
+
     def __repr__(self):
         return f"{type(self).__name__}: {self.rowCount()} rows"
 

prettyqt/custom_models/__init__.py

@@ -15,6 +15,7 @@
 from .jsonmodel import JsonModel
 from .subsequencesortfilterproxymodel import SubsequenceSortFilterProxyModel
 from .fuzzyfiltermodel import FuzzyFilterProxyModel
+from .subsetfilterproxymodel import SubsetFilterProxyModel
 
 __all__ = [
     "SelectionMixin",
@@ -31,4 +32,5 @@
     "JsonModel",
     "SubsequenceSortFilterProxyModel",
     "FuzzyFilterProxyModel",
+    "SubsetFilterProxyModel",
 ]

prettyqt/custom_models/subsetfilterproxymodel.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from collections.abc import Container
+
+from prettyqt import core
+from prettyqt.utils import get_repr
+
+
+def is_in_slice(a_slice: slice | range, idx: int) -> bool:
+    start = a_slice.start or 0
+    if idx < start or (a_slice.stop is not None and idx >= a_slice.stop):
+        return False
+    step = a_slice.step or 1
+    return (idx - start) % step == 0
+
+
+class SubsetFilterProxyModel(core.SortFilterProxyModel):
+    def __init__(
+        self,
+        row_filter: slice | range | int | Container[int] | None,
+        column_filter: slice | range | int | Container[int] | None,
+        **kwargs,
+    ):
+        self.row_filter = row_filter
+        self.column_filter = column_filter
+        super().__init__(**kwargs)
+
+    def __repr__(self):
+        return get_repr(self, self.row_filter, self.column_filter)
+
+    def filterAcceptsColumn(self, source_column: int, parent: core.ModelIndex) -> bool:
+        match self.column_filter:
+            case slice() | range():
+                return is_in_slice(self.column_filter, source_column)
+            case int():
+                return source_column == self.column_filter
+            case Container():
+                return source_column in self.column_filter
+            case None:
+                return True
+            case _:
+                raise ValueError(self.column_filter)
+
+    def filterAcceptsRow(self, source_row: int, parent: core.ModelIndex) -> bool:
+        match self.row_filter:
+            case slice() | range():
+                return is_in_slice(self.row_filter, source_row)
+            case int():
+                return source_row == self.row_filter
+            case Container():
+                return source_row in self.row_filter
+            case None:
+                return True
+            case _:
+                raise ValueError(self.row_filter)
+
+
+if __name__ == "__main__":
+    from prettyqt import widgets
+
+    app = widgets.app()
+    proxy = SubsetFilterProxyModel(1, 2)