initial commit

Author: PythonFZ

README.md

@@ -0,0 +1,91 @@
+# ZnSlice
+
+A lightweight library  (without external dependencies) for:
+- advanced slicing.
+- cache `__getitem__(self, item)`
+- lazy load `__getitem__(self, item)`
+
+# Installation
+
+```bash
+pip install znslice
+```
+
+# Usage
+
+## Advanced Slicing and Cache
+```python
+import znslice
+import collections.abc
+
+class MapList(collections.abc.Sequence):
+    def __init__(self, data, func):
+        self.data = data
+        self.func = func
+    
+    @znslice.znslice
+    def __getitem__(self, item: int):
+        print(f"Loading item = {item}")
+        return self.func(self.data[item])
+    
+    def __len__(self):
+        return len(self.data)
+
+data = MapList([0, 1, 2, 3, 4], lambda x: x ** 2)
+
+assert data[0] == 0
+assert data[[1, 2, 3]] == [1, 4, 9]
+# calling data[:] will now only compute data[4] and load the remaining data from cache
+assert data[:] == [0, 1, 4, 9, 16]
+```
+
+## Lazy Database Loading
+
+You can use `znslice` to lazy load data from a database. This is useful if you have a large database and only want to load a small subset of the data.
+
+In the following we will use the `ase` package to generate `Atoms` objects stored in a database and load them lazily.
+
+```python 
+import ase.io
+import ase.db
+import znslice
+import tqdm
+import random
+
+# create a database
+with ase.db.connect("data.db", append=False) as db:
+    for _ in range(10):
+        atoms = ase.Atoms('CO', positions=[(0, 0, 0), (0, 0, random.random())])
+        db.write(atoms, group="data")
+        
+# load the database lazily
+class ReadASEDB:
+    def __init__(self, file):
+        self.file = file
+    
+    @znslice.znslice(
+        advanced_slicing=True, # this getitem supports advanced slicingn
+        lazy=True # we want to lazy load the data
+    )
+    def __getitem__(self, item):
+        data = []
+        with ase.db.connect(self.file) as database:
+            if isinstance(item, int):
+                print(f"get {item = }")
+                return database[item + 1].toatoms()
+            for idx in tqdm.tqdm(item):
+                data.append(database[idx + 1].toatoms())
+        return data
+            
+    def __len__(self):
+        with ase.db.connect(self.file) as db:
+            return len(db)
+
+db = ReadASEDB("data.db")
+
+data = db[::2] # LazySequence([<__main__.ReadASEDB>], [[0, 2, 4, 6, 8]])
+data.tolist() # list[ase.Atoms] 
+
+# supports addition, advanced slicing, etc.
+data = db[::2] + db[1::2]
+```

poetry.lock

@@ -0,0 +1,27 @@
+[tool.poetry]
+name = "ZnSlice"
+version = "0.1.0"
+description = "Cache, advanced slicing and lazy loading for __getitem__"
+license = "Apache-2.0"
+authors = ["zincwarecode <[email protected]>"]
+readme = "README.md"
+
+[tool.poetry.urls]
+repository = "https://github.com/zincware/ZnSlice"
+
+[tool.poetry.dependencies]
+python = "^3.8"
+
+[tool.poetry.group.dev.dependencies]
+black = "^22.12.0"
+isort = "^5.11.4"
+ruff = "^0.0.224"
+pytest = "^7.2.1"
+ipykernel = "^6.20.2"
+ase = "^3.22.1"
+tqdm = "^4.64.1"
+coverage = "^7.0.5"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

pyproject.toml

@@ -0,0 +1,55 @@
+import collections.abc
+
+import znslice
+
+
+class DatabaseMock:
+    def __init__(self, maximum: int = 100):
+        self.maximum = maximum
+
+    def __getitem__(self, item):
+        if isinstance(item, int):
+            if item > self.maximum:
+                raise IndexError
+            return item
+        raise TypeError(
+            f"Index of type {type(item)} not supported. Only int supported."
+        )
+
+    def __len__(self):
+        return self.maximum
+
+
+class LazyList(collections.abc.Sequence):
+    """Lazy List for ASE Atoms Objects"""
+
+    def __init__(self, obj=None, indices=None):
+        self._obj = obj
+        self._indices = (
+            znslice.utils.item_to_indices(indices, range(len(obj))) if indices else None
+        )
+
+    @znslice.znslice(lazy=True)
+    def __getitem__(self, item: int):
+        if self._indices is None:
+            return self._obj[item]
+        return self._obj[self._indices[item]]
+
+    def __len__(self):
+        if self._indices is None:
+            return len(self._obj)
+        return len(self._indices)
+
+
+def test_LazyList():
+    lsta = LazyList(DatabaseMock(), indices=[1, 4, 7])
+    assert lsta[0] == 1
+    assert isinstance(lsta[:2], znslice.LazySequence)
+
+    assert lsta[[0, 1, 2]].tolist() == [1, 4, 7]
+
+    lstb = LazyList(DatabaseMock(), indices=[9, 11, 26])
+    assert lstb[[0, 1, 2]].tolist() == [9, 11, 26]
+
+    lstc = LazyList(DatabaseMock(), indices=slice(None, None, 3))
+    assert lstc[:3].tolist() == [0, 3, 6]

tests/test_LazyList.py

@@ -0,0 +1,16 @@
+import pytest
+
+import znslice
+
+
+def test_item_to_indices():
+
+    lst = list(range(10))
+
+    assert znslice.utils.item_to_indices(1, lst) == 1
+    assert znslice.utils.item_to_indices(slice(4, 6), lst) == [4, 5]
+    assert znslice.utils.item_to_indices([1, 2, 3], lst) == [1, 2, 3]
+    assert znslice.utils.item_to_indices((1, 2, 3), lst) == [1, 2, 3]
+
+    with pytest.raises(ValueError):
+        znslice.utils.item_to_indices("data", lst)