pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/exprs/expr_dict.py

@@ -1,8 +1,9 @@
-from typing import Generic, TypeVar, Optional, Iterator, Iterable
+from typing import Generic, Iterable, Iterator, TypeVar
+
+from .expr import Expr
 
 T = TypeVar('T')
 
-from .expr import Expr
 
 class ExprDict(Generic[T]):
     """
@@ -13,7 +14,7 @@ class ExprDict(Generic[T]):
 
     _data: dict[int, tuple[Expr, T]]
 
-    def __init__(self, iterable: Optional[Iterable[tuple[Expr, T]]] = None):
+    def __init__(self, iterable: Iterable[tuple[Expr, T]] | None = None):
         self._data = {}
 
         if iterable is not None:
@@ -38,7 +39,7 @@ class ExprDict(Generic[T]):
     def __contains__(self, key: Expr) -> bool:
         return key.id in self._data
 
-    def get(self, key: Expr, default: Optional[T] = None) -> Optional[T]:
+    def get(self, key: Expr, default: T | None = None) -> T | None:
         item = self._data.get(key.id)
         return item[1] if item is not None else default
 
@@ -46,7 +47,7 @@ class ExprDict(Generic[T]):
         self._data.clear()
 
     def keys(self) -> Iterator[Expr]:
-        return self.__iter__()
+        return iter(self)
 
     def values(self) -> Iterator[T]:
         return (value for _, value in self._data.values())

pixeltable/exprs/expr_set.py

@@ -1,32 +1,41 @@
 from __future__ import annotations
 
-from typing import Optional, Iterable, Iterator, TypeVar, Generic
+from typing import Generic, Iterable, Iterator, TypeVar
 
 from .expr import Expr
 
 T = TypeVar('T', bound='Expr')
 
+
 class ExprSet(Generic[T]):
     """
-    A set that also supports indexed lookup (by slot_idx and Expr.id). Exprs are uniquely identified by Expr.id.
+    An ordered set that also supports indexed lookup (by slot_idx and Expr.id). Exprs are uniquely identified by
+    Expr.id.
     """
+
     exprs: dict[int, T]  # key: Expr.id
+    expr_offsets: dict[int, int]  # key: Expr.id, value: offset into self.exprs.keys()
     exprs_by_idx: dict[int, T]  # key: slot_idx
 
-    def __init__(self, elements: Optional[Iterable[T]] = None):
+    def __init__(self, elements: Iterable[T] | None = None):
         self.exprs = {}
+        self.expr_offsets = {}
         self.exprs_by_idx = {}
         if elements is not None:
             for e in elements:
                 self.add(e)
 
-    def add(self, expr: T) -> None:
-        if expr.id in self.exprs:
-            return
+    def add(self, expr: T) -> int:
+        """Returns offset corresponding to iteration order"""
+        offset = self.expr_offsets.get(expr.id)
+        if offset is not None:
+            return offset
+        offset = len(self.exprs)
         self.exprs[expr.id] = expr
-        if expr.slot_idx is None:
-            return
-        self.exprs_by_idx[expr.slot_idx] = expr
+        self.expr_offsets[expr.id] = offset
+        if expr.slot_idx is not None:
+            self.exprs_by_idx[expr.slot_idx] = expr
+        return offset
 
     def update(self, *others: Iterable[T]) -> None:
         for other in others:
@@ -42,9 +51,9 @@ class ExprSet(Generic[T]):
     def __iter__(self) -> Iterator[T]:
         return iter(self.exprs.values())
 
-    def __getitem__(self, index: object) -> Optional[T]:
+    def __getitem__(self, index: object) -> T | None:
         """Indexed lookup by slot_idx or Expr.id."""
-        assert isinstance(index, int) or isinstance(index, Expr)
+        assert isinstance(index, (int, Expr))
         if isinstance(index, int):
             # return expr with matching slot_idx
             return self.exprs_by_idx.get(index)