pixeltable 0.2.17__py3-none-any.whl → 0.2.19__py3-none-any.whl

pixeltable/exprs/in_predicate.py

@@ -5,6 +5,7 @@ from typing import Optional, List, Any, Dict, Tuple, Iterable
 import sqlalchemy as sql
 
 import pixeltable.exceptions as excs
+from .sql_element_cache import SqlElementCache
 import pixeltable.type_system as ts
 from .data_row import DataRow
 from .expr import Expr
@@ -70,8 +71,8 @@ class InPredicate(Expr):
     def _id_attrs(self) -> List[Tuple[str, Any]]:
         return super()._id_attrs() + [('value_list', self.value_list)]
 
-    def sql_expr(self) -> Optional[sql.ClauseElement]:
-        lhs_sql_exprs = self.components[0].sql_expr()
+    def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ClauseElement]:
+        lhs_sql_exprs = sql_elements.get(self.components[0])
         if lhs_sql_exprs is None or self.value_list is None:
             return None
         return lhs_sql_exprs.in_(self.value_list)

pixeltable/exprs/inline_expr.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+import copy
+from typing import Any, Iterable, Optional
+
+import numpy as np
+import sqlalchemy as sql
+
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+
+from .data_row import DataRow
+from .expr import Expr
+from .literal import Literal
+from .row_builder import RowBuilder
+from .sql_element_cache import SqlElementCache
+
+
+class InlineArray(Expr):
+    """
+    Array 'literal' which can use Exprs as values.
+    """
+
+    def __init__(self, elements: Iterable):
+        exprs = []
+        for el in elements:
+            if isinstance(el, Expr):
+                exprs.append(el)
+            elif isinstance(el, list) or isinstance(el, tuple):
+                exprs.append(InlineArray(el))
+            else:
+                exprs.append(Literal(el))
+
+        inferred_element_type: Optional[ts.ColumnType] = ts.InvalidType()
+        for i, expr in enumerate(exprs):
+            supertype = inferred_element_type.supertype(expr.col_type)
+            if supertype is None:
+                raise excs.Error(
+                    f'Could not infer element type of array: element of type `{expr.col_type}` at index {i} '
+                    f'is not compatible with type `{inferred_element_type}` of preceding elements'
+                )
+            inferred_element_type = supertype
+
+        if inferred_element_type.is_scalar_type():
+            col_type = ts.ArrayType((len(exprs),), inferred_element_type)
+        elif inferred_element_type.is_array_type():
+            assert isinstance(inferred_element_type, ts.ArrayType)
+            col_type = ts.ArrayType(
+                (len(exprs), *inferred_element_type.shape),
+                ts.ColumnType.make_type(inferred_element_type.dtype)
+            )
+        else:
+            raise excs.Error(f'Element type is not a valid dtype for an array: {inferred_element_type}')
+
+        super().__init__(col_type)
+        self.components.extend(exprs)
+        self.id = self._create_id()
+
+    def __str__(self) -> str:
+        elem_strs = [str(expr) for expr in self.components]
+        return f'[{", ".join(elem_strs)}]'
+
+    def _equals(self, _: InlineArray) -> bool:
+        return True  # Always true if components match
+
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+        return None
+
+    def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
+        data_row[self.slot_idx] = np.array([data_row[el.slot_idx] for el in self.components])
+
+    def _as_dict(self) -> dict:
+        return super()._as_dict()
+
+    @classmethod
+    def _from_dict(cls, _: dict, components: list[Expr]) -> Expr:
+        try:
+            return cls(components)
+        except excs.Error:
+            # For legacy compatibility reasons, we need to try constructing as an `InlineList`.
+            # This is because in schema versions <= 19, `InlineArray` was serialized incorrectly, and
+            # there is no way to determine the correct expression type until the subexpressions are
+            # loaded and their types are known.
+            return InlineList(components)
+
+
+class InlineList(Expr):
+    """
+    List 'literal' which can use Exprs as values.
+    """
+
+    def __init__(self, elements: Iterable):
+        exprs = []
+        for el in elements:
+            if isinstance(el, Expr):
+                exprs.append(el)
+            elif isinstance(el, list) or isinstance(el, tuple):
+                exprs.append(InlineList(el))
+            elif isinstance(el, dict):
+                exprs.append(InlineDict(el))
+            else:
+                exprs.append(Literal(el))
+
+        super().__init__(ts.JsonType())
+        self.components.extend(exprs)
+        self.id = self._create_id()
+
+    def __str__(self) -> str:
+        elem_strs = [str(expr) for expr in self.components]
+        return f'[{", ".join(elem_strs)}]'
+
+    def _equals(self, _: InlineList) -> bool:
+        return True  # Always true if components match
+
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+        return None
+
+    def eval(self, data_row: DataRow, _: RowBuilder) -> None:
+        data_row[self.slot_idx] = [data_row[el.slot_idx] for el in self.components]
+
+    def _as_dict(self) -> dict:
+        return super()._as_dict()
+
+    @classmethod
+    def _from_dict(cls, _: dict, components: list[Expr]) -> Expr:
+        return cls(components)
+
+
+class InlineDict(Expr):
+    """
+    Dictionary 'literal' which can use Exprs as values.
+    """
+
+    keys: list[str]
+
+    def __init__(self, d: dict[str, Any]):
+        self.keys = []
+        exprs: list[Expr] = []
+        for key, val in d.items():
+            if not isinstance(key, str):
+                raise excs.Error(f'Dictionary requires string keys; {key} has type {type(key)}')
+            self.keys.append(key)
+            if isinstance(val, Expr):
+                exprs.append(val)
+            elif isinstance(val, dict):
+                exprs.append(InlineDict(val))
+            elif isinstance(val, list) or isinstance(val, tuple):
+                exprs.append(InlineList(val))
+            else:
+                exprs.append(Literal(val))
+
+        super().__init__(ts.JsonType())
+        self.components.extend(exprs)
+        self.id = self._create_id()
+
+    def __str__(self) -> str:
+        item_strs = list(f"'{key}': {str(expr)}" for key, expr in zip(self.keys, self.components))
+        return '{' + ', '.join(item_strs) + '}'
+
+    def _equals(self, other: InlineDict) -> bool:
+        # The dict values are just the components, which have already been checked
+        return self.keys == other.keys
+
+    def _id_attrs(self) -> list[tuple[str, Any]]:
+        return super()._id_attrs() + [('keys', self.keys)]
+
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+        return None
+
+    def eval(self, data_row: DataRow, _: RowBuilder) -> None:
+        assert len(self.keys) == len(self.components)
+        data_row[self.slot_idx] = {
+            key: data_row[expr.slot_idx]
+            for key, expr in zip(self.keys, self.components)
+        }
+
+    def to_kwargs(self) -> dict[str, Any]:
+        """Deconstructs this expression into a dictionary by recursively unwrapping all Literals,
+        InlineDicts, and InlineLists."""
+        return InlineDict._to_kwarg_element(self)
+
+    @classmethod
+    def _to_kwarg_element(cls, expr: Expr) -> Any:
+        if isinstance(expr, Literal):
+            return expr.val
+        if isinstance(expr, InlineDict):
+            return {key: cls._to_kwarg_element(val) for key, val in zip(expr.keys, expr.components)}
+        if isinstance(expr, InlineList):
+            return [cls._to_kwarg_element(el) for el in expr.components]
+        return expr
+
+    def _as_dict(self) -> dict[str, Any]:
+        return {'keys': self.keys, **super()._as_dict()}
+
+    @classmethod
+    def _from_dict(cls, d: dict, components: list[Expr]) -> Expr:
+        assert 'keys' in d
+        assert len(d['keys']) == len(components)
+        arg = dict(zip(d['keys'], components))
+        return InlineDict(arg)

pixeltable/exprs/is_null.py

@@ -8,6 +8,7 @@ import pixeltable.type_system as ts
 from .data_row import DataRow
 from .expr import Expr
 from .row_builder import RowBuilder
+from .sql_element_cache import SqlElementCache
 
 
 class IsNull(Expr):
@@ -22,8 +23,8 @@ class IsNull(Expr):
     def _equals(self, other: IsNull) -> bool:
         return True
 
-    def sql_expr(self) -> Optional[sql.ClauseElement]:
-        e = self.components[0].sql_expr()
+    def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ClauseElement]:
+        e = sql_elements.get(self.components[0])
         if e is None:
             return None
         return e == None

pixeltable/exprs/json_mapper.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
+
 from typing import Optional, List, Dict
 
 import sqlalchemy as sql
 
-from .expr import Expr, ExprScope, _GLOBAL_SCOPE
+import pixeltable.type_system as ts
 from .data_row import DataRow
+from .expr import Expr, ExprScope, _GLOBAL_SCOPE
 from .row_builder import RowBuilder
-import pixeltable.catalog as catalog
-import pixeltable.type_system as ts
+from .sql_element_cache import SqlElementCache
 
 
 class JsonMapper(Expr):
@@ -86,7 +87,7 @@ class JsonMapper(Expr):
     def _equals(self, other: JsonMapper) -> bool:
         return True
 
-    def sql_expr(self) -> Optional[sql.ClauseElement]:
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ClauseElement]:
         return None
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:

pixeltable/exprs/json_path.py

@@ -9,6 +9,12 @@ import pixeltable
 import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
+from .data_row import DataRow
+from .expr import Expr
+from .globals import print_slice
+from .json_mapper import JsonMapper
+from .row_builder import RowBuilder
+from .sql_element_cache import SqlElementCache
 
 from .data_row import DataRow
 from .expr import Expr
@@ -140,7 +146,7 @@ class JsonPath(Expr):
     def _id_attrs(self) -> list[tuple[str, Any]]:
         return super()._id_attrs() + [('path_elements', self.path_elements)]
 
-    def sql_expr(self) -> Optional[sql.ClauseElement]:
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ClauseElement]:
         """
         Postgres appears to have a bug: jsonb_path_query('{a: [{b: 0}, {b: 1}]}', '$.a.b') returns
         *two* rows (each containing col val 0), not a single row with [0, 0].

pixeltable/exprs/literal.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
 import datetime
-from typing import Optional, List, Any, Dict, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import sqlalchemy as sql
 
-import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
+from pixeltable.env import Env
+
 from .data_row import DataRow
 from .expr import Expr
 from .row_builder import RowBuilder
+from .sql_element_cache import SqlElementCache
 
 
 class Literal(Expr):
@@ -22,6 +24,15 @@ class Literal(Expr):
             if col_type is None:
                 raise TypeError(f'Not a valid literal: {val}')
         super().__init__(col_type)
+        if isinstance(val, datetime.datetime):
+            # Normalize the datetime to UTC: all timestamps are stored as UTC (both in the database and in literals)
+            if val.tzinfo is None:
+                # We have a naive datetime. Modify it to use the configured default time zone
+                default_tz = Env.get().default_time_zone
+                if default_tz is not None:
+                    val = val.replace(tzinfo=default_tz)
+            # Now convert to UTC
+            val = val.astimezone(datetime.timezone.utc)
         self.val = val
         self.id = self._create_id()
 
@@ -29,17 +40,24 @@ class Literal(Expr):
         return 'Literal'
 
     def __str__(self) -> str:
-        if self.col_type.is_string_type() or self.col_type.is_timestamp_type():
+        if self.col_type.is_string_type():
             return f"'{self.val}'"
+        if self.col_type.is_timestamp_type():
+            assert isinstance(self.val, datetime.datetime)
+            default_tz = Env.get().default_time_zone
+            return f"'{self.val.astimezone(default_tz).isoformat()}'"
         return str(self.val)
 
+    def __repr__(self) -> str:
+        return f'Literal({self.val!r})'
+
     def _equals(self, other: Literal) -> bool:
         return self.val == other.val
 
     def _id_attrs(self) -> List[Tuple[str, Any]]:
         return super()._id_attrs() + [('val', self.val)]
 
-    def sql_expr(self) -> Optional[sql.ClauseElement]:
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ClauseElement]:
         # we need to return something here so that we can generate a Where clause for predicates
         # that involve literals (like Where c > 0)
         return sql.sql.expression.literal(self.val)
@@ -52,7 +70,12 @@ class Literal(Expr):
         # For some types, we need to explictly record their type, because JSON does not know
         # how to interpret them unambiguously
         if self.col_type.is_timestamp_type():
-            return {'val': self.val.isoformat(), 'val_t': self.col_type._type.name, **super()._as_dict()}
+            assert isinstance(self.val, datetime.datetime)
+            assert self.val.tzinfo == datetime.timezone.utc  # Must be UTC in a literal
+            # Convert to ISO format in UTC (in keeping with the principle: all timestamps are
+            # stored as UTC in the database)
+            encoded_val = self.val.isoformat()
+            return {'val': encoded_val, 'val_t': self.col_type._type.name, **super()._as_dict()}
         else:
             return {'val': self.val, **super()._as_dict()}
 
@@ -61,6 +84,10 @@ class Literal(Expr):
         assert 'val' in d
         if 'val_t' in d:
             val_t = d['val_t']
+            # Currently the only special-cased literal type is TIMESTAMP
             assert val_t == ts.ColumnType.Type.TIMESTAMP.name
-            return cls(datetime.datetime.fromisoformat(d['val']))
-        return cls(d['val'])
+            dt = datetime.datetime.fromisoformat(d['val'])
+            assert dt.tzinfo == datetime.timezone.utc  # Must be UTC in the database
+            return cls(dt)
+        else:
+            return cls(d['val'])

pixeltable/exprs/method_ref.py

@@ -2,12 +2,12 @@ from typing import Any, Optional
 
 import sqlalchemy as sql
 
-import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 from pixeltable.exprs import Expr, FunctionCall
-from pixeltable.func import FunctionRegistry, CallableFunction
+from pixeltable.func import FunctionRegistry
 from .data_row import DataRow
 from .row_builder import RowBuilder
+from .sql_element_cache import SqlElementCache
 
 
 class MethodRef(Expr):
@@ -53,7 +53,7 @@ class MethodRef(Expr):
     def _id_attrs(self) -> list[tuple[str, Any]]:
         return super()._id_attrs() + [('method_name', self.method_name)]
 
-    def sql_expr(self) -> Optional[sql.ClauseElement]:
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ClauseElement]:
         return None
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:

pixeltable/exprs/object_ref.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
-from typing import Optional, List, Any, Dict, Tuple
-import copy
+
+from typing import Optional
 
 import sqlalchemy as sql
 
+import pixeltable.type_system as ts
+from .data_row import DataRow
 from .expr import Expr, ExprScope
 from .json_mapper import JsonMapper
-from .data_row import DataRow
 from .row_builder import RowBuilder
-import pixeltable.type_system as ts
+from .sql_element_cache import SqlElementCache
 
 
 class ObjectRef(Expr):
@@ -32,7 +33,7 @@ class ObjectRef(Expr):
     def _equals(self, other: ObjectRef) -> bool:
         return self.owner is other.owner
 
-    def sql_expr(self) -> Optional[sql.ClauseElement]:
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ClauseElement]:
         return None
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:

pixeltable/exprs/row_builder.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import sys
 import time
 from dataclasses import dataclass
-from typing import Optional, List, Any, Dict, Sequence, Tuple, Set
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple
 
 import sqlalchemy as sql
 
@@ -11,6 +11,7 @@ import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
 import pixeltable.func as func
 import pixeltable.utils as utils
+
 from .data_row import DataRow
 from .expr import Expr
 from .expr_set import ExprSet
@@ -22,7 +23,7 @@ class ExecProfile:
         self.eval_count = [0] * row_builder.num_materialized
         self.row_builder = row_builder
 
-    def print(self, num_rows: int) -> str:
+    def print(self, num_rows: int) -> None:
         for i in range(self.row_builder.num_materialized):
             if self.eval_count[i] == 0:
                 continue
@@ -57,7 +58,7 @@ class RowBuilder:
         target_exprs: List[Expr]  # exprs corresponding to target_slot_idxs
 
     def __init__(
-            self, output_exprs: Sequence[Expr], columns: Sequence[catalog.Column], input_exprs: Sequence[Expr]
+            self, output_exprs: Sequence[Expr], columns: Sequence[catalog.Column], input_exprs: Iterable[Expr]
     ):
         """
         Args:
@@ -96,7 +97,7 @@ class RowBuilder:
                 expr = ColumnRef(col)
                 expr = self._record_unique_expr(expr, recursive=False)
             self.add_table_column(col, expr.slot_idx)
-            self.output_exprs.append(expr)
+            self.output_exprs.add(expr)
 
         # default eval ctx: all output exprs
         self.default_eval_ctx = self.create_eval_ctx(list(self.output_exprs), exclude=unique_input_exprs)
@@ -193,7 +194,7 @@ class RowBuilder:
                 expr.components[i] = self._record_unique_expr(c, True)
         assert expr.slot_idx is None
         expr.slot_idx = self._next_slot_idx()
-        self.unique_exprs.append(expr)
+        self.unique_exprs.add(expr)
         return expr
 
     def _record_output_expr_id(self, e: Expr, output_expr_id: int) -> None:
@@ -227,18 +228,25 @@ class RowBuilder:
         # merge dependencies and convert to list
         return sorted(set().union(*[dependencies[i] for i in target_slot_idxs]))
 
-    def substitute_exprs(self, expr_list: list, remove_duplicates: bool = True) -> None:
-        """Substitutes exprs with their executable counterparts from unique_exprs and optionally removes duplicates"""
-        i = 0
-        unique_ids: set[int] = set()  # slot idxs within expr_list
-        while i < len(expr_list):
-            unique_expr = self.unique_exprs[expr_list[i]]
-            if unique_expr.slot_idx in unique_ids and remove_duplicates:
-                del expr_list[i]
-            else:
-                expr_list[i] = unique_expr
-                unique_ids.add(unique_expr.slot_idx)
-                i += 1
+    def set_slot_idxs(self, expr_list: Sequence[Expr], remove_duplicates: bool = True) -> None:
+        """
+        Recursively sets slot_idx in expr_list and its components
+
+        remove_duplicates == True: removes duplicates in-place
+        """
+        for e in expr_list:
+            self.__set_slot_idxs_aux(e)
+        if remove_duplicates:
+            deduped = list(ExprSet(expr_list))
+            expr_list[:] = deduped
+
+    def __set_slot_idxs_aux(self, e: Expr) -> None:
+        """Recursively sets slot_idx in e and its components"""
+        if e not in self.unique_exprs:
+            return
+        e.slot_idx = self.unique_exprs[e].slot_idx
+        for c in e.components:
+            self.__set_slot_idxs_aux(c)
 
     def get_dependencies(self, targets: List[Expr], exclude: Optional[List[Expr]] = None) -> List[Expr]:
         """

pixeltable/exprs/rowid_ref.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 from typing import Optional, List, Any, Dict, Tuple
+from .sql_element_cache import SqlElementCache
 from uuid import UUID
 
 import sqlalchemy as sql
@@ -72,7 +73,7 @@ class RowidRef(Expr):
         self.tbl = tbl.tbl_version
         self.tbl_id = self.tbl.id
 
-    def sql_expr(self) -> Optional[sql.ClauseElement]:
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ClauseElement]:
         tbl = self.tbl if self.tbl is not None else catalog.Catalog.get().tbl_versions[(self.tbl_id, None)]
         rowid_cols = tbl.store_tbl.rowid_columns()
         return rowid_cols[self.rowid_component_idx]

pixeltable/exprs/similarity_expr.py

@@ -1,4 +1,5 @@
 from typing import Optional, List, Any
+from .sql_element_cache import SqlElementCache
 
 import sqlalchemy as sql
 import PIL.Image
@@ -56,7 +57,7 @@ class SimilarityExpr(Expr):
     def __str__(self) -> str:
         return f'{self.components[0]}.similarity({self.components[1]})'
 
-    def sql_expr(self) -> Optional[sql.ClauseElement]:
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ClauseElement]:
         if not isinstance(self.components[1], Literal):
              raise excs.Error(f'similarity(): requires a string or a PIL.Image.Image object, not an expression')
         item = self.components[1].val

pixeltable/exprs/sql_element_cache.py

@@ -0,0 +1,30 @@
+from typing import Iterable, Union, Optional
+
+import sqlalchemy as sql
+
+from .expr import Expr
+
+
+class SqlElementCache:
+    """Cache of sql.ColumnElements for exprs"""
+
+    cache: dict[int, Optional[sql.ColumnElement]]  # key: Expr.id
+
+    def __init__(self):
+        self.cache = {}
+
+    def get(self, e: Expr) -> Optional[sql.ColumnElement]:
+        """Returns the sql.ColumnElement for the given Expr, or None if Expr.to_sql() returns None."""
+        try:
+            return self.cache[e.id]
+        except KeyError:
+            pass
+        el = e.sql_expr(self)
+        self.cache[e.id] = el
+        return el
+
+    def contains(self, items: Union[Expr, Iterable[Expr]]) -> bool:
+        """Returns True if every item has a (non-None) sql.ColumnElement."""
+        if isinstance(items, Expr):
+            return self.get(items) is not None
+        return all(self.get(e) is not None for e in items)

pixeltable/exprs/type_cast.py

@@ -1,4 +1,3 @@
-import json
 from typing import Optional, Dict, List, Tuple, Any
 
 import sqlalchemy as sql
@@ -6,6 +5,7 @@ import sqlalchemy as sql
 import pixeltable.type_system as ts
 from .expr import DataRow, Expr
 from .row_builder import RowBuilder
+from .sql_element_cache import SqlElementCache
 
 
 class TypeCast(Expr):
@@ -29,9 +29,9 @@ class TypeCast(Expr):
     def _id_attrs(self) -> List[Tuple[str, Any]]:
         return super()._id_attrs() + [('new_type', self.col_type)]
 
-    def sql_expr(self) -> Optional[sql.ClauseElement]:
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ClauseElement]:
         """
-        `sql_expr` is unimplemented for now, in order to sidestep potentially thorny
+        sql_expr() is unimplemented for now, in order to sidestep potentially thorny
         questions about consistency of doing type conversions in both Python and Postgres.
         """
         return None

pixeltable/exprs/variable.py

@@ -6,6 +6,7 @@ import pixeltable.type_system as ts
 from .data_row import DataRow
 from .expr import Expr
 from .row_builder import RowBuilder
+from .sql_element_cache import SqlElementCache
 
 
 class Variable(Expr):
@@ -31,7 +32,7 @@ class Variable(Expr):
     def __str__(self) -> str:
         return self.name
 
-    def sql_expr(self) -> NoReturn:
+    def sql_expr(self, _: SqlElementCache) -> NoReturn:
         raise NotImplementedError()
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> NoReturn:

pixeltable/ext/functions/whisperx.py

@@ -1,9 +1,9 @@
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
-    from whisperx.asr import FasterWhisperPipeline
+    from whisperx.asr import FasterWhisperPipeline  # type: ignore[import-untyped]
 
 import pixeltable as pxt
 
@@ -19,6 +19,8 @@ def transcribe(
     equivalent to the WhisperX `transcribe` function, as described in the
     [WhisperX library documentation](https://github.com/m-bain/whisperX).
 
+    WhisperX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
+
     __Requirements:__
 
     - `pip install whisperx`
@@ -40,7 +42,7 @@ def transcribe(
         >>> tbl['result'] = transcribe(tbl.audio, model='tiny.en')
     """
     import torch
-    import whisperx
+    import whisperx  # type: ignore[import-untyped]
 
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     compute_type = compute_type or ('float16' if device == 'cuda' else 'int8')
@@ -60,7 +62,7 @@ def _lookup_model(model_id: str, device: str, compute_type: str) -> 'FasterWhisp
     return _model_cache[key]
 
 
-_model_cache = {}
+_model_cache: dict[tuple[str, str, str], 'FasterWhisperPipeline'] = {}
 
 
 __all__ = local_public_names(__name__)