pixeltable 0.3.10__py3-none-any.whl → 0.3.12__py3-none-any.whl

pixeltable/exprs/json_mapper.py

@@ -20,16 +20,84 @@ class JsonMapper(Expr):
     JsonMapper transforms the list output of a JsonPath by applying a target expr to every element of the list.
     The target expr would typically contain relative JsonPaths, which are bound to an ObjectRef, which in turn
     is populated by JsonMapper.eval(). The JsonMapper effectively creates a new scope for its target expr.
+
+    JsonMapper is executed in two phases:
+    - the first phase is handled by Expr subclass JsonMapperDispatch, which constructs one nested DataRow per source
+      list element and evaluates the target expr within that (the nested DataRows are stored as a NestedRowList in the
+      slot of JsonMapperDispatch)
+    - JsonMapper.eval() collects the slot values of the target expr into its result list
     """
 
     target_expr_scope: ExprScope
     parent_mapper: Optional[JsonMapper]
     target_expr_eval_ctx: Optional[RowBuilder.EvalCtx]
 
-    def __init__(self, src_expr: Expr, target_expr: Expr):
+    def __init__(self, src_expr: Optional[Expr], target_expr: Optional[Expr]):
         # TODO: type spec should be list[target_expr.col_type]
         super().__init__(ts.JsonType())
 
+        dispatch = JsonMapperDispatch(src_expr, target_expr)
+        self.components.append(dispatch)
+        self.id = self._create_id()
+
+    def __repr__(self) -> str:
+        return f'map({self._src_expr}, lambda R: {self._target_expr})'
+
+    @property
+    def _src_expr(self) -> Expr:
+        return self.components[0].src_expr
+
+    @property
+    def _target_expr(self) -> Expr:
+        return self.components[0].target_expr
+
+    def _equals(self, _: JsonMapper) -> bool:
+        return True
+
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+        return None
+
+    def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
+        from ..exec.expr_eval.evaluators import NestedRowList
+
+        dispatch_slot_idx = self.components[0].slot_idx
+        nested_rows = data_row.vals[dispatch_slot_idx]
+        if nested_rows is None:
+            data_row[self.slot_idx] = None
+            return
+        assert isinstance(nested_rows, NestedRowList)
+        # TODO: get the materialized slot idx, instead of relying on the fact that the target_expr is always at the end
+        data_row[self.slot_idx] = [row.vals[-1] for row in nested_rows.rows]
+
+    def _as_dict(self) -> dict:
+        """
+        We only serialize src and target exprs, everything else is re-created at runtime.
+        """
+        return {'components': [self._src_expr.as_dict(), self._target_expr.as_dict()]}
+
+    @classmethod
+    def _from_dict(cls, d: dict, components: list[Expr]) -> JsonMapper:
+        assert len(components) == 2
+        src_expr, target_expr = components[0], components[1]
+        return cls(src_expr, target_expr)
+
+
+class JsonMapperDispatch(Expr):
+    """
+    An operational Expr (ie, it doesn't represent any syntactic element) that is used by JsonMapper to materialize
+    its input DataRows. It has the same dependencies as the originating JsonMapper.
+
+    - The execution (= row dispatch) is handled by an expr_eval.Evaluator (JsonMapperDispatcher).
+    - It stores a NestedRowList instance in its slot.
+    """
+
+    target_expr_scope: ExprScope
+    parent_mapper: Optional[JsonMapperDispatch]
+    target_expr_eval_ctx: Optional[RowBuilder.EvalCtx]
+
+    def __init__(self, src_expr: Expr, target_expr: Expr):
+        super().__init__(ts.InvalidType())
+
         # we're creating a new scope, but we don't know yet whether this is nested within another JsonMapper;
         # this gets resolved in bind_rel_paths(); for now we assume we're in the global scope
         self.target_expr_scope = ExprScope(_GLOBAL_SCOPE)
@@ -40,28 +108,36 @@ class JsonMapper(Expr):
         self.parent_mapper = None
         self.target_expr_eval_ctx = None
 
-        # Intentionally create the id now, before adding the scope anchor; this ensures that JsonMappers will
-        # be recognized as equal so long as they have the same src_expr and target_expr.
+        # Intentionally create the id now, before adding the scope anchor; this ensures that JsonMapperDispatch
+        # instances will be recognized as equal so long as they have the same src_expr and target_expr.
         # TODO: Might this cause problems after certain substitutions?
         self.id = self._create_id()
 
         scope_anchor = ObjectRef(self.target_expr_scope, self)
         self.components.append(scope_anchor)
 
-    def _bind_rel_paths(self, mapper: Optional[JsonMapper] = None) -> None:
-        self._src_expr._bind_rel_paths(mapper)
-        self._target_expr._bind_rel_paths(self)
+    def _bind_rel_paths(self, mapper: Optional[JsonMapperDispatch] = None) -> None:
+        self.src_expr._bind_rel_paths(mapper)
+        self.target_expr._bind_rel_paths(self)
         self.parent_mapper = mapper
         parent_scope = _GLOBAL_SCOPE if mapper is None else mapper.target_expr_scope
         self.target_expr_scope.parent = parent_scope
 
+    def equals(self, other: Expr) -> bool:
+        """
+        We override equals() because we need to avoid comparing our scope anchor.
+        """
+        if type(self) is not type(other):
+            return False
+        return self.src_expr.equals(other.src_expr) and self.target_expr.equals(other.target_expr)
+
     def scope(self) -> ExprScope:
         # need to ignore target_expr
-        return self._src_expr.scope()
+        return self.src_expr.scope()
 
     def dependencies(self) -> list[Expr]:
-        result = [self._src_expr]
-        result.extend(self._target_dependencies(self._target_expr))
+        result = [self.src_expr]
+        result.extend(self._target_dependencies(self.target_expr))
         return result
 
     def _target_dependencies(self, e: Expr) -> list[Expr]:
@@ -77,23 +153,12 @@ class JsonMapper(Expr):
             result.extend(self._target_dependencies(c))
         return result
 
-    def equals(self, other: Expr) -> bool:
-        """
-        We override equals() because we need to avoid comparing our scope anchor.
-        """
-        if type(self) is not type(other):
-            return False
-        return self._src_expr.equals(other._src_expr) and self._target_expr.equals(other._target_expr)
-
-    def __repr__(self) -> str:
-        return f'map({self._src_expr}, lambda R: {self._target_expr})'
-
     @property
-    def _src_expr(self) -> Expr:
+    def src_expr(self) -> Expr:
         return self.components[0]
 
     @property
-    def _target_expr(self) -> Expr:
+    def target_expr(self) -> Expr:
         return self.components[1]
 
     @property
@@ -104,37 +169,19 @@ class JsonMapper(Expr):
         assert isinstance(result, ObjectRef)
         return result
 
-    def _equals(self, _: JsonMapper) -> bool:
-        return True
-
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
-        return None
+    def __repr__(self) -> str:
+        return 'JsonMapperDispatch()'
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
-        # this will be called, but the value has already been materialized elsewhere
-        src = data_row[self._src_expr.slot_idx]
-        if not isinstance(src, list):
-            # invalid/non-list src path
-            data_row[self.slot_idx] = None
-            return
-
-        result = [None] * len(src)
-        if self.target_expr_eval_ctx is None:
-            self.target_expr_eval_ctx = row_builder.create_eval_ctx([self._target_expr])
-        for i, val in enumerate(src):
-            data_row[self.scope_anchor.slot_idx] = val
-            # stored target_expr
-            row_builder.eval(data_row, self.target_expr_eval_ctx, force_eval=self._target_expr.scope())
-            result[i] = data_row[self._target_expr.slot_idx]
-        data_row[self.slot_idx] = result
+        # eval is handled by JsonMapperDispatcher
+        raise AssertionError('this should never be called')
 
     def _as_dict(self) -> dict:
         """
-        We need to avoid serializing component[2], which is an ObjectRef.
+        JsonMapperDispatch instances are only created by the JsonMapper c'tor and never need to be serialized.
         """
-        return {'components': [c.as_dict() for c in self.components[0:2]]}
+        raise AssertionError('this should never be called')
 
     @classmethod
-    def _from_dict(cls, d: dict, components: list[Expr]) -> JsonMapper:
-        assert len(components) == 2
-        return cls(components[0], components[1])
+    def _from_dict(cls, d: dict, components: list[Expr]) -> JsonMapperDispatch:
+        raise AssertionError('this should never be called')

pixeltable/exprs/json_path.py

@@ -5,13 +5,12 @@ from typing import Any, Optional, Union
 import jmespath
 import sqlalchemy as sql
 
-import pixeltable as pxt
 from pixeltable import catalog, exceptions as excs, type_system as ts
 
 from .data_row import DataRow
 from .expr import Expr
 from .globals import print_slice
-from .json_mapper import JsonMapper
+from .json_mapper import JsonMapperDispatch
 from .object_ref import ObjectRef
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
@@ -19,10 +18,7 @@ from .sql_element_cache import SqlElementCache
 
 class JsonPath(Expr):
     def __init__(
-        self,
-        anchor: Optional['pxt.exprs.Expr'],
-        path_elements: Optional[list[Union[str, int, slice]]] = None,
-        scope_idx: int = 0,
+        self, anchor: Optional[Expr], path_elements: Optional[list[Union[str, int, slice]]] = None, scope_idx: int = 0
     ) -> None:
         """
         anchor can be None, in which case this is a relative JsonPath and the anchor is set later via set_anchor().
@@ -80,11 +76,10 @@ class JsonPath(Expr):
     def is_relative_path(self) -> bool:
         return self._anchor is None
 
-    @property
     def _has_relative_path(self) -> bool:
-        return self.is_relative_path() or super()._has_relative_path
+        return self.is_relative_path() or super()._has_relative_path()
 
-    def _bind_rel_paths(self, mapper: Optional['JsonMapper'] = None) -> None:
+    def _bind_rel_paths(self, mapper: Optional['JsonMapperDispatch'] = None) -> None:
         if self.is_relative_path():
             # TODO: take scope_idx into account
             self.set_anchor(mapper.scope_anchor)

pixeltable/exprs/method_ref.py

@@ -23,7 +23,7 @@ class MethodRef(Expr):
     # TODO: Should this even be an `Expr`? It can't actually be evaluated directly (it has to be first
     #   converted to a `FunctionCall` by binding any remaining parameters).
 
-    def __init__(self, base_expr: Expr, method_name: str):
+    def __init__(self, base_expr: Expr, method_name: str) -> None:
         super().__init__(ts.InvalidType())  # The `MethodRef` is untyped until it is called.
         self.base_expr = base_expr
         self.method_name = method_name
@@ -43,7 +43,7 @@ class MethodRef(Expr):
         assert len(components) == 1
         return cls(components[0], d['method_name'])
 
-    def __call__(self, *args, **kwargs) -> FunctionCall:
+    def __call__(self, *args: Any, **kwargs: Any) -> FunctionCall:
         result = self.fn(*[self.base_expr, *args], **kwargs)
         assert isinstance(result, FunctionCall)
         result.is_method_call = True

pixeltable/exprs/object_ref.py

@@ -8,7 +8,7 @@ import pixeltable.type_system as ts
 
 from .data_row import DataRow
 from .expr import Expr, ExprScope
-from .json_mapper import JsonMapper
+from .json_mapper import JsonMapperDispatch
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
 
@@ -19,7 +19,7 @@ class ObjectRef(Expr):
     The object is generated/materialized elsewhere and establishes a new scope.
     """
 
-    def __init__(self, scope: ExprScope, owner: JsonMapper):
+    def __init__(self, scope: ExprScope, owner: JsonMapperDispatch):
         # TODO: do we need an Unknown type after all?
         super().__init__(ts.JsonType())  # JsonType: this could be anything
         self._scope = scope

pixeltable/exprs/row_builder.py

@@ -77,6 +77,8 @@ class RowBuilder:
     transitive_dependents: np.ndarray  # of bool
     # dependencies[i] = direct dependencies of expr with slot idx i; transpose of dependents
     dependencies: np.ndarray  # of bool
+    # num_dependencies[i] = number of direct dependencies of expr with slot idx i
+    num_dependencies: np.ndarray  # of int
 
     # records the output_expr that a subexpr belongs to
     # (a subexpr can be shared across multiple output exprs)
@@ -209,6 +211,7 @@ class RowBuilder:
                 exc_dependencies[expr.slot_idx].add(d.slot_idx)
                 exc_dependencies[expr.slot_idx].update(exc_dependencies[d.slot_idx])
 
+        self.num_dependencies = np.sum(self.dependencies, axis=1)
         self.dependents = self.dependencies.T
         self.transitive_dependents = np.zeros((self.num_materialized, self.num_materialized), dtype=bool)
         for i in reversed(range(self.num_materialized)):
@@ -275,8 +278,14 @@ class RowBuilder:
         for d in e.dependencies():
             self._record_output_expr_id(d, output_expr_id)
 
-    def _compute_dependencies(self, target_slot_idxs: list[int], excluded_slot_idxs: list[int]) -> list[int]:
-        """Compute exprs needed to materialize the given target slots, excluding 'excluded_slot_idxs'"""
+    def _compute_dependencies(
+        self, target_slot_idxs: list[int], excluded_slot_idxs: list[int], target_scope: Optional[ExprScope] = None
+    ) -> list[int]:
+        """Compute exprs needed to materialize the given target slots, excluding 'excluded_slot_idxs'
+
+        If target_scope != None, stops transitive dependency resolution when leaving target_scope (ie, includes
+        immediate dependents that aren't in target_scope, but doesn't resolve those).
+        """
         dependencies: list[set[int]] = [set() for _ in range(self.num_materialized)]  # indexed by slot_idx
         # doing this front-to-back ensures that we capture transitive dependencies
         max_target_slot_idx = max(target_slot_idxs)
@@ -289,6 +298,9 @@ class RowBuilder:
             if expr.slot_idx in self.input_expr_slot_idxs:
                 # this is input and therefore doesn't depend on other exprs
                 continue
+            if target_scope is not None and expr.scope() != target_scope:
+                # don't resolve dependencies outside of target_scope
+                continue
             for d in expr.dependencies():
                 assert d.slot_idx is not None, f'{expr}, {d}'
                 if d.slot_idx in excluded_slot_idxs:
@@ -320,10 +332,15 @@ class RowBuilder:
         for c in e.components:
             self.__set_slot_idxs_aux(c)
 
-    def get_dependencies(self, targets: Iterable[Expr], exclude: Optional[Iterable[Expr]] = None) -> list[Expr]:
+    def get_dependencies(
+        self, targets: Iterable[Expr], exclude: Optional[Iterable[Expr]] = None, limit_scope: bool = True
+    ) -> list[Expr]:
         """
         Return list of dependencies needed to evaluate the given target exprs (expressed as slot idxs).
         The exprs given in 'exclude' are excluded.
+        If limit_scope == True, only returns dependencies in the same scope and immediate (ie, not transitive)
+        dependencies from enclosing scopes.
+
         Returns:
             list of Exprs from unique_exprs (= with slot_idx set)
         """
@@ -334,23 +351,33 @@ class RowBuilder:
             return []
         # make sure we only refer to recorded exprs
         targets = [self.unique_exprs[e] for e in targets]
+        target_scope: Optional[ExprScope] = None
+        if limit_scope:
+            # make sure all targets are from the same scope
+            target_scopes = {e.scope() for e in targets}
+            assert len(target_scopes) == 1
+            target_scope = target_scopes.pop()
         exclude = [self.unique_exprs[e] for e in exclude]
         target_slot_idxs = [e.slot_idx for e in targets]
         excluded_slot_idxs = [e.slot_idx for e in exclude]
-        all_dependencies = set(self._compute_dependencies(target_slot_idxs, excluded_slot_idxs))
+        all_dependencies = set(
+            self._compute_dependencies(target_slot_idxs, excluded_slot_idxs, target_scope=target_scope)
+        )
         all_dependencies.update(target_slot_idxs)
         result_ids = list(all_dependencies)
         result_ids.sort()
         return [self.unique_exprs[id] for id in result_ids]
 
-    def create_eval_ctx(self, targets: Iterable[Expr], exclude: Optional[Iterable[Expr]] = None) -> EvalCtx:
+    def create_eval_ctx(
+        self, targets: Iterable[Expr], exclude: Optional[Iterable[Expr]] = None, limit_scope: bool = True
+    ) -> EvalCtx:
         """Return EvalCtx for targets"""
         targets = list(targets)
         if exclude is None:
             exclude = []
         if len(targets) == 0:
             return self.EvalCtx([], [], [], [])
-        dependencies = self.get_dependencies(targets, exclude)
+        dependencies = self.get_dependencies(targets, exclude, limit_scope=limit_scope)
         targets = [self.unique_exprs[e] for e in targets]
         target_slot_idxs = [e.slot_idx for e in targets]
         ctx_slot_idxs = [e.slot_idx for e in dependencies]

pixeltable/exprs/similarity_expr.py

@@ -47,7 +47,7 @@ class SimilarityExpr(Expr):
     def __repr__(self) -> str:
         return f'{self.components[0]}.similarity({self.components[1]})'
 
-    def _id_attrs(self):
+    def _id_attrs(self) -> list[tuple[str, Any]]:
         return [*super()._id_attrs(), ('idx_name', self.idx_info.name)]
 
     def default_column_name(self) -> str:

pixeltable/exprs/sql_element_cache.py

@@ -17,7 +17,7 @@ class SqlElementCache:
             for e, el in elements.items():
                 self.cache[e.id] = el
 
-    def extend(self, elements: ExprDict[sql.ColumnElement]):
+    def extend(self, elements: ExprDict[sql.ColumnElement]) -> None:
         for e, el in elements.items():
             self.cache[e.id] = el
 

pixeltable/exprs/string_op.py

@@ -26,7 +26,7 @@ class StringOp(Expr):
         self.operator = operator
         self.components = [op1, op2]
         assert op1.col_type.is_string_type()
-        if operator in {StringOperator.CONCAT, StringOperator.REPEAT}:
+        if operator in (StringOperator.CONCAT, StringOperator.REPEAT):
             if operator == StringOperator.CONCAT and not op2.col_type.is_string_type():
                 raise excs.Error(
                     f'{self}: {operator} on strings requires string type, but {op2} has type {op2.col_type}'
@@ -89,7 +89,7 @@ class StringOp(Expr):
         """
         Return the result of evaluating the expression on two int/float operands
         """
-        assert self.operator in {StringOperator.CONCAT, StringOperator.REPEAT}
+        assert self.operator in (StringOperator.CONCAT, StringOperator.REPEAT)
         if self.operator == StringOperator.CONCAT:
             assert isinstance(op2_val, str)
             return op1_val + op2_val

pixeltable/ext/__init__.py

@@ -13,5 +13,5 @@ from . import functions
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/ext/functions/__init__.py

@@ -7,5 +7,5 @@ from . import whisperx, yolox
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/ext/functions/whisperx.py

@@ -73,5 +73,5 @@ _model_cache: dict[tuple[str, str, str], 'FasterWhisperPipeline'] = {}
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/ext/functions/yolox.py

@@ -110,5 +110,5 @@ _processor_cache: dict[str, 'YoloxProcessor'] = {}
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/func/__init__.py

@@ -5,7 +5,7 @@ from .callable_function import CallableFunction
 from .expr_template_function import ExprTemplateFunction
 from .function import Function, InvalidFunction
 from .function_registry import FunctionRegistry
-from .query_template_function import QueryTemplateFunction, query
+from .query_template_function import QueryTemplateFunction, query, retrieval_udf
 from .signature import Batch, Parameter, Signature
 from .tools import Tool, ToolChoice, Tools
 from .udf import expr_udf, make_function, udf

pixeltable/func/aggregate_function.py

@@ -159,7 +159,7 @@ class AggregateFunction(Function):
         self.init_param_names.append(init_param_names)
         return self
 
-    def _docstring(self) -> Optional[str]:
+    def comment(self) -> Optional[str]:
         return inspect.getdoc(self.agg_classes[0])
 
     def help_str(self) -> str:
@@ -252,7 +252,7 @@ def uda(
 ) -> Callable[[type[Aggregator]], AggregateFunction]: ...
 
 
-def uda(*args, **kwargs):
+def uda(*args, **kwargs):  # type: ignore[no-untyped-def]
     """Decorator for user-defined aggregate functions.
 
     The decorated class must inherit from Aggregator and implement the following methods:

pixeltable/func/callable_function.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import asyncio
 import inspect
 from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence
 from uuid import UUID
@@ -61,7 +60,7 @@ class CallableFunction(Function):
     def is_async(self) -> bool:
         return inspect.iscoroutinefunction(self.py_fn)
 
-    def _docstring(self) -> Optional[str]:
+    def comment(self) -> Optional[str]:
         return inspect.getdoc(self.py_fns[0])
 
     @property
@@ -127,12 +126,10 @@ class CallableFunction(Function):
         """
         assert self.is_batched
         assert not self.is_polymorphic
+        assert not self.is_async
         # Unpack the constant parameters
         constant_kwargs, batched_kwargs = self.create_batch_kwargs(kwargs)
-        if inspect.iscoroutinefunction(self.py_fn):
-            return asyncio.run(self.py_fn(*args, **constant_kwargs, **batched_kwargs))
-        else:
-            return self.py_fn(*args, **constant_kwargs, **batched_kwargs)
+        return self.py_fn(*args, **constant_kwargs, **batched_kwargs)
 
     def create_batch_kwargs(self, kwargs: dict[str, Any]) -> tuple[dict[str, Any], dict[str, list[Any]]]:
         """Converts kwargs containing lists into constant and batched kwargs in the format expected by a batched udf."""

pixeltable/func/expr_template_function.py

@@ -1,6 +1,6 @@
 from typing import Any, Optional, Sequence
 
-from pixeltable import exceptions as excs, exprs
+from pixeltable import exceptions as excs, exprs, type_system as ts
 
 from .function import Function
 from .signature import Signature
@@ -76,12 +76,28 @@ class ExprTemplateFunction(Function):
                 arg_expr = arg
             arg_exprs[param_expr] = arg_expr
         result = result.substitute(arg_exprs)
-        assert not result._contains(exprs.Variable)
         return result
 
-    def _docstring(self) -> Optional[str]:
+    def call_return_type(self, bound_args: dict[str, 'exprs.Expr']) -> ts.ColumnType:
+        """
+        The call_return_type of an ExprTemplateFunction is derived from the template expression's col_type after
+        substitution (unlike for UDFs, whose call_return_type is derived from an explicitly specified
+        conditional_return_type).
+        """
+        assert not self.is_polymorphic
+        template = self.template
+        with_defaults = bound_args.copy()
+        with_defaults.update(
+            {param_name: default for param_name, default in template.defaults.items() if param_name not in bound_args}
+        )
+        substituted_expr = self.template.expr.copy().substitute(
+            {template.param_exprs[name]: expr for name, expr in with_defaults.items()}
+        )
+        return substituted_expr.col_type
+
+    def comment(self) -> Optional[str]:
         if isinstance(self.templates[0].expr, exprs.FunctionCall):
-            return self.templates[0].expr.fn._docstring()
+            return self.templates[0].expr.fn.comment()
         return None
 
     def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:
@@ -97,6 +113,10 @@ class ExprTemplateFunction(Function):
 
     @property
     def display_name(self) -> str:
+        if not self.self_name and isinstance(self.templates[0].expr, exprs.FunctionCall):
+            # In the common case where the templated expression is itself a FunctionCall,
+            # fall back on the display name of the underlying FunctionCall
+            return self.templates[0].expr.fn.display_name
         return self.self_name
 
     @property

pixeltable/func/function.py

@@ -10,9 +10,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, cast
 import sqlalchemy as sql
 from typing_extensions import Self
 
-import pixeltable as pxt
-import pixeltable.exceptions as excs
-import pixeltable.type_system as ts
+from pixeltable import exceptions as excs, type_system as ts
 
 from .globals import resolve_symbol
 from .signature import Signature
@@ -107,11 +105,11 @@ class Function(ABC):
     @abstractmethod
     def is_async(self) -> bool: ...
 
-    def _docstring(self) -> Optional[str]:
+    def comment(self) -> Optional[str]:
         return None
 
     def help_str(self) -> str:
-        docstring = self._docstring()
+        docstring = self.comment()
         display = self.display_name + str(self.signatures[0])
         if docstring is None:
             return display
@@ -155,7 +153,7 @@ class Function(ABC):
         """
         raise NotImplementedError()
 
-    def __call__(self, *args: Any, **kwargs: Any) -> 'pxt.exprs.FunctionCall':
+    def __call__(self, *args: Any, **kwargs: Any) -> 'exprs.FunctionCall':
         from pixeltable import exprs
 
         args = [exprs.Expr.from_object(arg) for arg in args]
@@ -246,7 +244,7 @@ class Function(ABC):
         # `None` when any of its non-nullable inputs are `None`.
         for arg_name, arg in bound_args.items():
             param = self.signature.parameters[arg_name]
-            if param.kind in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}:
+            if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
                 continue
             if arg.col_type.nullable and not param.col_type.nullable:
                 return_type = return_type.copy(nullable=True)
@@ -385,10 +383,10 @@ class Function(ABC):
             else:
                 var = exprs.Variable(name, param.col_type)
                 bindings[name] = var
-                if args_ok and param.kind in {
+                if args_ok and param.kind in (
                     inspect.Parameter.POSITIONAL_ONLY,
                     inspect.Parameter.POSITIONAL_OR_KEYWORD,
-                }:
+                ):
                     template_args.append(var)
                 else:
                     template_kwargs[name] = var

pixeltable/func/function_registry.py

@@ -31,7 +31,7 @@ class FunctionRegistry:
             cls._instance = FunctionRegistry()
         return cls._instance
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.stored_fns_by_id: dict[UUID, Function] = {}
         self.module_fns: dict[str, Function] = {}  # fqn -> Function
         self.type_methods: dict[ts.ColumnType.Type, dict[str, Function]] = {}