pixeltable 0.3.7__py3-none-any.whl → 0.3.9__py3-none-any.whl

pixeltable/catalog/column.py

@@ -132,6 +132,7 @@ class Column:
             from pixeltable import exprs
 
             self._value_expr = exprs.Expr.from_dict(self.value_expr_dict)
+            self._value_expr.bind_rel_paths()
             if not self._value_expr.is_valid:
                 message = (
                     dedent(

pixeltable/catalog/dir.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
 import dataclasses
+import datetime
+import json
 import logging
 from uuid import UUID
 
 import sqlalchemy as sql
+from sqlalchemy.dialects.postgresql import JSONB
 
 from pixeltable.env import Env
 from pixeltable.metadata import schema
@@ -26,6 +29,7 @@ class Dir(SchemaObject):
         dir_record = schema.Dir(parent_id=parent_id, md=dataclasses.asdict(dir_md))
         session.add(dir_record)
         session.flush()
+        # print(f'{datetime.datetime.now()} create dir {dir_record}')
         assert dir_record.id is not None
         assert isinstance(dir_record.id, UUID)
         dir = cls(dir_record.id, parent_id, name)
@@ -43,11 +47,16 @@ class Dir(SchemaObject):
         return super()._path()
 
     def _move(self, new_name: str, new_dir_id: UUID) -> None:
+        # print(
+        #     f'{datetime.datetime.now()} move dir name={self._name} parent={self._dir_id} new_name={new_name} new_dir_id={new_dir_id}'
+        # )
         super()._move(new_name, new_dir_id)
-        with Env.get().engine.begin() as conn:
-            dir_md = schema.DirMd(name=new_name, user=None, additional_md={})
-            conn.execute(
-                sql.update(schema.Dir.__table__)
-                .values({schema.Dir.parent_id: self._dir_id, schema.Dir.md: dataclasses.asdict(dir_md)})
-                .where(schema.Dir.id == self._id)
+        stmt = sql.text(
+            (
+                f'UPDATE {schema.Dir.__table__} '
+                f'SET {schema.Dir.parent_id.name} = :new_dir_id, '
+                f"    {schema.Dir.md.name}['name'] = :new_name "
+                f'WHERE {schema.Dir.id.name} = :id'
             )
+        )
+        Env.get().conn.execute(stmt, {'new_dir_id': new_dir_id, 'new_name': json.dumps(new_name), 'id': self._id})

pixeltable/catalog/path.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from typing import Iterator
 
 from pixeltable import exceptions as excs
 
@@ -55,5 +56,19 @@ class Path:
         is_prefix = self.components == other.components[: self.len]
         return is_prefix and (self.len == (other.len - 1) or not is_parent)
 
+    def ancestors(self) -> Iterator[Path]:
+        """
+        Return all ancestors of this path in top-down order including root.
+        If this path is for the root directory, which has no parent, then None is returned.
+        """
+        if self.is_root:
+            return
+        else:
+            for i in range(0, len(self.components)):
+                yield Path('.'.join(self.components[0:i]), empty_is_valid=True)
+
     def __str__(self) -> str:
         return '.'.join(self.components)
+
+    def __lt__(self, other: Path) -> bool:
+        return str(self) < str(other)

pixeltable/catalog/schema_object.py

@@ -2,7 +2,7 @@ from abc import abstractmethod
 from typing import TYPE_CHECKING, Any, Optional
 from uuid import UUID
 
-import pixeltable.env as env
+from pixeltable.env import Env
 
 if TYPE_CHECKING:
     from pixeltable import catalog
@@ -28,24 +28,19 @@ class SchemaObject:
         """Returns the parent directory of this schema object."""
         from .catalog import Catalog
 
-        with env.Env.get().begin_xact():
+        with Env.get().begin_xact():
             if self._dir_id is None:
                 return None
             return Catalog.get().get_dir(self._dir_id)
 
     def _path(self) -> str:
         """Returns the path to this schema object."""
-        with env.Env.get().begin_xact():
-            from .catalog import Catalog
+        from .catalog import Catalog
 
-            cat = Catalog.get()
-            dir_path = cat.get_dir_path(self._dir_id)
-            if dir_path == '':
-                # Either this is the root directory, with empty path, or its parent is the
-                # root directory. Either way, we return just the name.
-                return self._name
-            else:
-                return f'{dir_path}.{self._name}'
+        assert self._dir_id is not None
+        with Env.get().begin_xact():
+            path = Catalog.get().get_dir_path(self._dir_id)
+            return str(path.append(self._name))
 
     def get_metadata(self) -> dict[str, Any]:
         """Returns metadata associated with this schema object."""

pixeltable/catalog/table.py

@@ -171,8 +171,8 @@ class Table(SchemaObject):
 
     def _get_views(self, *, recursive: bool = True) -> list['Table']:
         cat = catalog.Catalog.get()
-        view_ids = cat.get_views(self._id)
-        views = [cat.get_tbl(id) for id in view_ids]
+        view_ids = cat.get_view_ids(self._id)
+        views = [cat.get_table_by_id(id) for id in view_ids]
         if recursive:
             views.extend([t for view in views for t in view._get_views(recursive=True)])
         return views
@@ -265,7 +265,7 @@ class Table(SchemaObject):
         if self._tbl_version_path.base is None:
             return None
         base_id = self._tbl_version_path.base.tbl_version.id
-        return catalog.Catalog.get().get_tbl(base_id)
+        return catalog.Catalog.get().get_table_by_id(base_id)
 
     @property
     def _bases(self) -> list['Table']:
@@ -369,11 +369,6 @@ class Table(SchemaObject):
             pd_rows.append(row)
         return pd.DataFrame(pd_rows)
 
-    def ensure_md_loaded(self) -> None:
-        """Ensure that table metadata is loaded."""
-        for col in self._tbl_version.get().cols_by_id.values():
-            _ = col.value_expr
-
     def describe(self) -> None:
         """
         Print the table schema.
@@ -387,13 +382,9 @@ class Table(SchemaObject):
             print(repr(self))
 
     def _drop(self) -> None:
-        cat = catalog.Catalog.get()
         self._check_is_dropped()
         self._tbl_version.get().drop()
         self._is_dropped = True
-        # update catalog
-        cat = catalog.Catalog.get()
-        cat.remove_tbl(self._id)
 
     # TODO Factor this out into a separate module.
     # The return type is unresolvable, but torch can't be imported since it's an optional dependency.

pixeltable/catalog/table_version.py

@@ -454,6 +454,11 @@ class TableVersion:
                 )
             )
 
+    def ensure_md_loaded(self) -> None:
+        """Ensure that table metadata is loaded."""
+        for col in self.cols_by_id.values():
+            _ = col.value_expr
+
     def _store_idx_name(self, idx_id: int) -> str:
         """Return name of index in the store, which needs to be globally unique"""
         return f'idx_{self.id.hex}_{idx_id}'

pixeltable/catalog/view.py

@@ -237,15 +237,11 @@ class View(Table):
         )
 
     def _drop(self) -> None:
-        cat = catalog.Catalog.get()
         if self._snapshot_only:
             # there is not TableVersion to drop
             self._check_is_dropped()
             self.is_dropped = True
             TableVersion.delete_md(self._id)
-            # update catalog
-            cat = catalog.Catalog.get()
-            cat.remove_tbl(self._id)
         else:
             super()._drop()
 

pixeltable/env.py

@@ -170,19 +170,25 @@ class Env:
         assert self._current_session is not None
         return self._current_session
 
+    def in_xact(self) -> bool:
+        return self._current_conn is not None
+
     @contextmanager
     def begin_xact(self) -> Iterator[sql.Connection]:
         """Return a context manager that yields a connection to the database. Idempotent."""
         if self._current_conn is None:
             assert self._current_session is None
-            with self.engine.begin() as conn, sql.orm.Session(conn) as session:
-                self._current_conn = conn
-                self._current_session = session
-                try:
+            try:
+                with self.engine.begin() as conn, sql.orm.Session(conn) as session:
+                    # TODO: remove print() once we're done with debugging the concurrent update behavior
+                    # print(f'{datetime.datetime.now()}: start xact')
+                    self._current_conn = conn
+                    self._current_session = session
                     yield conn
-                finally:
-                    self._current_session = None
-                    self._current_conn = None
+            finally:
+                self._current_session = None
+                self._current_conn = None
+                # print(f'{datetime.datetime.now()}: end xact')
         else:
             assert self._current_session is not None
             yield self._current_conn
@@ -391,7 +397,7 @@ class Env:
     def _create_engine(self, time_zone_name: Optional[str], echo: bool = False) -> None:
         connect_args = {} if time_zone_name is None else {'options': f'-c timezone={time_zone_name}'}
         self._sa_engine = sql.create_engine(
-            self.db_url, echo=echo, future=True, isolation_level='REPEATABLE READ', connect_args=connect_args
+            self.db_url, echo=echo, isolation_level='REPEATABLE READ', connect_args=connect_args
         )
         self._logger.info(f'Created SQLAlchemy engine at: {self.db_url}')
         with self.engine.begin() as conn:

pixeltable/exprs/__init__.py

@@ -1,3 +1,5 @@
+# ruff: noqa: F401
+
 from .arithmetic_expr import ArithmeticExpr
 from .array_slice import ArraySlice
 from .column_property_ref import ColumnPropertyRef

pixeltable/exprs/arithmetic_expr.py

@@ -1,12 +1,10 @@
 from __future__ import annotations
 
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 import sqlalchemy as sql
 
-import pixeltable.exceptions as excs
-import pixeltable.exprs as exprs
-import pixeltable.type_system as ts
+from pixeltable import exceptions as excs, type_system as ts
 
 from .data_row import DataRow
 from .expr import Expr
@@ -50,13 +48,13 @@ class ArithmeticExpr(Expr):
         # add parentheses around operands that are ArithmeticExprs to express precedence
         op1_str = f'({self._op1})' if isinstance(self._op1, ArithmeticExpr) else str(self._op1)
         op2_str = f'({self._op2})' if isinstance(self._op2, ArithmeticExpr) else str(self._op2)
-        return f'{op1_str} {str(self.operator)} {op2_str}'
+        return f'{op1_str} {self.operator} {op2_str}'
 
     def _equals(self, other: ArithmeticExpr) -> bool:
         return self.operator == other.operator
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('operator', self.operator.value)]
+        return [*super()._id_attrs(), ('operator', self.operator.value)]
 
     def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ColumnElement]:
         assert self.col_type.is_int_type() or self.col_type.is_float_type() or self.col_type.is_json_type()
@@ -95,7 +93,7 @@ class ArithmeticExpr(Expr):
                 return sql.sql.expression.cast(sql.func.floor(left / nullif), self.col_type.to_sa_type())
             if self.col_type.is_float_type():
                 return sql.sql.expression.cast(sql.func.floor(left / nullif), self.col_type.to_sa_type())
-        assert False
+        raise AssertionError()
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
         op1_val = data_row[self._op1.slot_idx]
@@ -113,9 +111,7 @@ class ArithmeticExpr(Expr):
 
         data_row[self.slot_idx] = self.eval_nullable(op1_val, op2_val)
 
-    def eval_nullable(
-        self, op1_val: Union[int, float, None], op2_val: Union[int, float, None]
-    ) -> Union[int, float, None]:
+    def eval_nullable(self, op1_val: Optional[float], op2_val: Optional[float]) -> Optional[float]:
         """
         Return the result of evaluating the expression on two nullable int/float operands,
         None is interpreted as SQL NULL
@@ -124,7 +120,7 @@ class ArithmeticExpr(Expr):
             return None
         return self.eval_non_null(op1_val, op2_val)
 
-    def eval_non_null(self, op1_val: Union[int, float], op2_val: Union[int, float]) -> Union[int, float]:
+    def eval_non_null(self, op1_val: float, op2_val: float) -> float:
         """
         Return the result of evaluating the expression on two int/float operands
         """

pixeltable/exprs/array_slice.py

@@ -41,7 +41,7 @@ class ArraySlice(Expr):
         return self.index == other.index
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('index', self.index)]
+        return [*super()._id_attrs(), ('index', self.index)]
 
     def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
         return None

pixeltable/exprs/column_property_ref.py

@@ -40,7 +40,7 @@ class ColumnPropertyRef(Expr):
         return self.prop == other.prop
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('prop', self.prop.value)]
+        return [*super()._id_attrs(), ('prop', self.prop.value)]
 
     @property
     def _col_ref(self) -> ColumnRef:
@@ -52,7 +52,7 @@ class ColumnPropertyRef(Expr):
         return f'{self._col_ref}.{self.prop.name.lower()}'
 
     def is_error_prop(self) -> bool:
-        return self.prop == self.Property.ERRORTYPE or self.prop == self.Property.ERRORMSG
+        return self.prop in {self.Property.ERRORTYPE, self.Property.ERRORMSG}
 
     def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ColumnElement]:
         if not self._col_ref.col.is_stored:
@@ -95,7 +95,7 @@ class ColumnPropertyRef(Expr):
             else:
                 data_row[self.slot_idx] = str(exc)
         else:
-            assert False
+            raise AssertionError()
 
     def _as_dict(self) -> dict:
         return {'prop': self.prop.value, **super()._as_dict()}

pixeltable/exprs/column_ref.py

@@ -6,9 +6,7 @@ from uuid import UUID
 import sqlalchemy as sql
 
 import pixeltable as pxt
-import pixeltable.catalog as catalog
-import pixeltable.exceptions as excs
-import pixeltable.iterators as iters
+from pixeltable import catalog, exceptions as excs, iterators as iters
 
 from ..utils.description_helper import DescriptionHelper
 from .data_row import DataRow
@@ -84,7 +82,8 @@ class ColumnRef(Expr):
         assert len(self.iter_arg_ctx.target_slot_idxs) == 1  # a single inline dict
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [
+        return [
+            *super()._id_attrs(),
             ('tbl_id', self.col.tbl.id),
             ('col_id', self.col.id),
             ('perform_validation', self.perform_validation),
@@ -138,7 +137,7 @@ class ColumnRef(Expr):
         return self.col == other.col and self.perform_validation == other.perform_validation
 
     def _df(self) -> 'pxt.dataframe.DataFrame':
-        tbl = catalog.Catalog.get().get_tbl(self.col.tbl.id)
+        tbl = catalog.Catalog.get().get_table_by_id(self.col.tbl.id)
         return tbl.select(self)
 
     def show(self, *args, **kwargs) -> 'pxt.dataframe.DataFrameResultSet':
@@ -166,7 +165,7 @@ class ColumnRef(Expr):
         return self._descriptors().to_html()
 
     def _descriptors(self) -> DescriptionHelper:
-        tbl = catalog.Catalog.get().get_tbl(self.col.tbl.id)
+        tbl = catalog.Catalog.get().get_table_by_id(self.col.tbl.id)
         helper = DescriptionHelper()
         helper.append(f'Column\n{self.col.name!r}\n(of table {tbl._path()!r})')
         helper.append(tbl._col_descriptor([self.col.name]))

pixeltable/exprs/comparison.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Optional
+from typing import Any, Optional
 
 import sqlalchemy as sql
 
@@ -15,9 +15,6 @@ from .literal import Literal
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
 
-if TYPE_CHECKING:
-    from pixeltable import index
-
 
 class Comparison(Expr):
     is_search_arg_comparison: bool
@@ -62,7 +59,7 @@ class Comparison(Expr):
         return self.operator == other.operator
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('operator', self.operator.value)]
+        return [*super()._id_attrs(), ('operator', self.operator.value)]
 
     @property
     def _op1(self) -> Expr:

pixeltable/exprs/compound_predicate.py

@@ -5,7 +5,7 @@ from typing import Any, Callable, Optional
 
 import sqlalchemy as sql
 
-import pixeltable.type_system as ts
+from pixeltable import type_system as ts
 
 from .data_row import DataRow
 from .expr import Expr
@@ -58,10 +58,10 @@ class CompoundPredicate(Expr):
         return self.operator == other.operator
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('operator', self.operator.value)]
+        return [*super()._id_attrs(), ('operator', self.operator.value)]
 
     def split_conjuncts(self, condition: Callable[[Expr], bool]) -> tuple[list[Expr], Optional[Expr]]:
-        if self.operator == LogicalOperator.OR or self.operator == LogicalOperator.NOT:
+        if self.operator in {LogicalOperator.OR, LogicalOperator.NOT}:
             return super().split_conjuncts(condition)
         matches = [op for op in self.components if condition(op)]
         non_matches = [op for op in self.components if not condition(op)]
@@ -83,7 +83,7 @@ class CompoundPredicate(Expr):
         if self.operator == LogicalOperator.NOT:
             data_row[self.slot_idx] = not data_row[self.components[0].slot_idx]
         else:
-            val = True if self.operator == LogicalOperator.AND else False
+            val = self.operator == LogicalOperator.AND
             op_function = operator.and_ if self.operator == LogicalOperator.AND else operator.or_
             for op in self.components:
                 val = op_function(val, data_row[op.slot_idx])

pixeltable/exprs/expr.py

@@ -90,14 +90,29 @@ class Expr(abc.ABC):
                 result = c_scope
         return result
 
-    def bind_rel_paths(self, mapper: Optional['exprs.JsonMapper'] = None) -> None:
+    def bind_rel_paths(self) -> None:
         """
         Binds relative JsonPaths to mapper.
         This needs to be done in a separate phase after __init__(), because RelativeJsonPath()(-1) cannot be resolved
         by the immediately containing JsonMapper during initialization.
         """
+        self._bind_rel_paths()
+        assert not self._has_relative_path, self._expr_tree()
+
+    def _bind_rel_paths(self, mapper: Optional['exprs.JsonMapper'] = None) -> None:
+        for c in self.components:
+            c._bind_rel_paths(mapper)
+
+    def _expr_tree(self) -> str:
+        """Returns a string representation of this expression as a multi-line tree. Useful for debugging."""
+        buf: list[str] = []
+        self._expr_tree_r(0, buf)
+        return '\n'.join(buf)
+
+    def _expr_tree_r(self, indent: int, buf: list[str]) -> None:
+        buf.append(f'{" " * indent}{type(self).__name__}: {self}'.replace('\n', '\\n'))
         for c in self.components:
-            c.bind_rel_paths(mapper)
+            c._expr_tree_r(indent + 2, buf)
 
     def default_column_name(self) -> Optional[str]:
         """
@@ -129,7 +144,7 @@ class Expr(abc.ABC):
         """
         Subclass-specific comparison. Implemented as a function because __eq__() is needed to construct Comparisons.
         """
-        if type(self) != type(other):
+        if type(self) is not type(other):
             return False
         if len(self.components) != len(other.components):
             return False
@@ -171,10 +186,7 @@ class Expr(abc.ABC):
     def list_equals(cls, a: list[Expr], b: list[Expr]) -> bool:
         if len(a) != len(b):
             return False
-        for i in range(len(a)):
-            if not a[i].equals(b[i]):
-                return False
-        return True
+        return all(a[i].equals(b[i]) for i in range(len(a)))
 
     def copy(self) -> Expr:
         """
@@ -216,9 +228,9 @@ class Expr(abc.ABC):
                 return new.copy()
         for i in range(len(self.components)):
             self.components[i] = self.components[i].substitute(spec)
-        self = self.maybe_literal()
-        self.id = self._create_id()
-        return self
+        result = self.maybe_literal()
+        result.id = result._create_id()
+        return result
 
     @classmethod
     def list_substitute(cls, expr_list: list[Expr], spec: dict[Expr, Expr]) -> None:
@@ -253,10 +265,7 @@ class Expr(abc.ABC):
         from .column_ref import ColumnRef
 
         col_refs = self.subexprs(ColumnRef)
-        for col_ref in col_refs:
-            if not any(tbl.has_column(col_ref.col) for tbl in tbls):
-                return False
-        return True
+        return all(any(tbl.has_column(col_ref.col) for tbl in tbls) for col_ref in col_refs)
 
     def retarget(self, tbl: catalog.TableVersionPath) -> Self:
         """Retarget ColumnRefs in this expr to the specific TableVersions in tbl."""
@@ -361,6 +370,10 @@ class Expr(abc.ABC):
         except StopIteration:
             return False
 
+    @property
+    def _has_relative_path(self) -> bool:
+        return any(c._has_relative_path for c in self.components)
+
     def tbl_ids(self) -> set[UUID]:
         """Returns table ids referenced by this expr."""
         from .column_ref import ColumnRef
@@ -370,7 +383,7 @@ class Expr(abc.ABC):
 
     @classmethod
     def all_tbl_ids(cls, exprs_: Iterable[Expr]) -> set[UUID]:
-        return set(tbl_id for e in exprs_ for tbl_id in e.tbl_ids())
+        return {tbl_id for e in exprs_ for tbl_id in e.tbl_ids()}
 
     @classmethod
     def get_refd_columns(cls, expr_dict: dict[str, Any]) -> list[catalog.Column]:
@@ -489,7 +502,7 @@ class Expr(abc.ABC):
         return {'_classname': self.__class__.__name__, **self._as_dict()}
 
     @classmethod
-    def as_dict_list(self, expr_list: list[Expr]) -> list[dict]:
+    def as_dict_list(cls, expr_list: list[Expr]) -> list[dict]:
         return [e.as_dict() for e in expr_list]
 
     def _as_dict(self) -> dict:
@@ -520,7 +533,7 @@ class Expr(abc.ABC):
 
     @classmethod
     def _from_dict(cls, d: dict, components: list[Expr]) -> Self:
-        assert False, 'not implemented'
+        raise AssertionError(f'not implemented: {cls.__name__}')
 
     def isin(self, value_set: Any) -> 'exprs.InPredicate':
         from .in_predicate import InPredicate
@@ -792,13 +805,13 @@ class Expr(abc.ABC):
             first_param = next(params_iter) if len(params) >= 1 else None
             second_param = next(params_iter) if len(params) >= 2 else None
             # Check that fn has at least one positional parameter
-            if len(params) == 0 or first_param.kind in (inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.VAR_KEYWORD):
+            if len(params) == 0 or first_param.kind in {inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.VAR_KEYWORD}:
                 raise excs.Error(f'Function `{fn.__name__}` has no positional parameters.')
             # Check that fn has at most one required parameter, i.e., its second parameter
             # has no default and is not a varargs
             if (
                 len(params) >= 2
-                and second_param.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
+                and second_param.kind not in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}
                 and second_param.default is inspect.Parameter.empty
             ):
                 raise excs.Error(f'Function `{fn.__name__}` has multiple required parameters.')

pixeltable/exprs/expr_dict.py

@@ -1,9 +1,9 @@
 from typing import Generic, Iterable, Iterator, Optional, TypeVar
 
-T = TypeVar('T')
-
 from .expr import Expr
 
+T = TypeVar('T')
+
 
 class ExprDict(Generic[T]):
     """
@@ -47,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

@@ -46,7 +46,7 @@ class ExprSet(Generic[T]):
 
     def __getitem__(self, index: object) -> Optional[T]:
         """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)