pixeltable 0.4.18__py3-none-any.whl → 0.4.19__py3-none-any.whl

pixeltable/exprs/similarity_expr.py

@@ -1,4 +1,6 @@
-from typing import Any, Optional
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 import sqlalchemy as sql
 
@@ -12,34 +14,43 @@ from .literal import Literal
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
 
+if TYPE_CHECKING:
+    from pixeltable.catalog.table_version import TableVersion
+
 
 class SimilarityExpr(Expr):
-    def __init__(self, col_ref: ColumnRef, item: Any, idx_name: Optional[str] = None):
+    """
+    A similarity expression against an embedding index.
+    """
+
+    idx_id: int
+    idx_name: str
+
+    def __init__(self, col_ref: ColumnRef, item: Any, idx_name: str | None = None):
+        from pixeltable.index import EmbeddingIndex
+
         super().__init__(ts.FloatType())
         item_expr = Expr.from_object(item)
         if item_expr is None or not (item_expr.col_type.is_string_type() or item_expr.col_type.is_image_type()):
             raise excs.Error(f'similarity(): requires a string or a PIL.Image.Image object, not a {type(item)}')
-        assert item_expr.col_type.is_string_type() or item_expr.col_type.is_image_type()
 
         self.components = [col_ref, item_expr]
 
-        from pixeltable import index
-
         # determine index to use
-        idx_dict = col_ref.find_embedding_index(idx_name, 'similarity')
-        assert len(idx_dict) == 1
-        self.idx_info = next(iter(idx_dict.values()))
-        idx = self.idx_info.idx
-        assert isinstance(idx, index.EmbeddingIndex)
+        idx_info = col_ref.tbl.get().get_idx(col_ref.col, idx_name, EmbeddingIndex)
+        self.idx_id = idx_info.id
+        self.idx_name = idx_info.name
+        idx = idx_info.idx
+        assert isinstance(idx, EmbeddingIndex)
 
         if item_expr.col_type.is_string_type() and idx.string_embed is None:
             raise excs.Error(
-                f'Embedding index {self.idx_info.name!r} on column {self.idx_info.col.name!r} does not have a '
+                f'Embedding index {idx_info.name!r} on column {idx_info.col.name!r} does not have a '
                 f'string embedding and does not support string queries'
             )
         if item_expr.col_type.is_image_type() and idx.image_embed is None:
             raise excs.Error(
-                f'Embedding index {self.idx_info.name!r} on column {self.idx_info.col.name!r} does not have an '
+                f'Embedding index {idx_info.name!r} on column {idx_info.col.name!r} does not have an '
                 f'image embedding and does not support image queries'
             )
         self.id = self._create_id()
@@ -48,39 +59,53 @@ class SimilarityExpr(Expr):
         return f'{self.components[0]}.similarity({self.components[1]})'
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return [*super()._id_attrs(), ('idx_name', self.idx_info.name)]
+        return [*super()._id_attrs(), ('idx_id', self.idx_id)]
 
     def default_column_name(self) -> str:
         return 'similarity'
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
-        # TODO: validate that the index still exists
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
+        from pixeltable.index import EmbeddingIndex
+
+        # check for a literal here, instead of the c'tor: needed for ExprTemplateFunctions
         if not isinstance(self.components[1], Literal):
             raise excs.Error('similarity(): requires a string or a PIL.Image.Image object, not an expression')
         item = self.components[1].val
-        from pixeltable import index
+        idx_info = self._resolve_idx()
+        assert isinstance(idx_info.idx, EmbeddingIndex)
+        return idx_info.idx.similarity_clause(idx_info.val_col, item)
 
-        assert isinstance(self.idx_info.idx, index.EmbeddingIndex)
-        return self.idx_info.idx.similarity_clause(self.idx_info.val_col, item)
+    def as_order_by_clause(self, is_asc: bool) -> sql.ColumnElement | None:
+        from pixeltable.index import EmbeddingIndex
 
-    def as_order_by_clause(self, is_asc: bool) -> Optional[sql.ColumnElement]:
+        # check for a literal here, instead of the c'tor: needed for ExprTemplateFunctions
         if not isinstance(self.components[1], Literal):
             raise excs.Error('similarity(): requires a string or a PIL.Image.Image object, not an expression')
         item = self.components[1].val
-        from pixeltable import index
+        idx_info = self._resolve_idx()
+        assert isinstance(idx_info.idx, EmbeddingIndex)
+        return idx_info.idx.order_by_clause(idx_info.val_col, item, is_asc)
+
+    def _resolve_idx(self) -> 'TableVersion.IndexInfo':
+        from pixeltable.index import EmbeddingIndex
 
-        assert isinstance(self.idx_info.idx, index.EmbeddingIndex)
-        return self.idx_info.idx.order_by_clause(self.idx_info.val_col, item, is_asc)
+        # resolve idx_id
+        col_ref = self.components[0]
+        if self.idx_id not in col_ref.tbl.get().idxs:
+            raise excs.Error(f'Index {self.idx_name!r} not found')
+        idx_info = col_ref.tbl.get().idxs[self.idx_id]
+        assert isinstance(idx_info.idx, EmbeddingIndex)
+        return idx_info
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
         raise excs.Error('similarity(): cannot be used in a computed column')
 
     def _as_dict(self) -> dict:
-        return {'idx_name': self.idx_info.name, **super()._as_dict()}
+        return {'idx_name': self.idx_name, **super()._as_dict()}
 
     @classmethod
     def _from_dict(cls, d: dict, components: list[Expr]) -> 'SimilarityExpr':
-        iname = d.get('idx_name')
+        idx_name = d.get('idx_name')
         assert len(components) == 2
         assert isinstance(components[0], ColumnRef)
-        return cls(components[0], components[1], idx_name=iname)
+        return cls(components[0], components[1], idx_name=idx_name)

pixeltable/exprs/sql_element_cache.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Optional
+from typing import Iterable
 
 import sqlalchemy as sql
 
@@ -9,9 +9,9 @@ from .expr_dict import ExprDict
 class SqlElementCache:
     """Cache of sql.ColumnElements for exprs"""
 
-    cache: dict[int, Optional[sql.ColumnElement]]  # key: Expr.id
+    cache: dict[int, sql.ColumnElement | None]  # key: Expr.id
 
-    def __init__(self, elements: Optional[ExprDict[sql.ColumnElement]] = None):
+    def __init__(self, elements: ExprDict[sql.ColumnElement] | None = None):
         self.cache = {}
         if elements is not None:
             for e, el in elements.items():
@@ -21,7 +21,7 @@ class SqlElementCache:
         for e, el in elements.items():
             self.cache[e.id] = el
 
-    def get(self, e: Expr) -> Optional[sql.ColumnElement]:
+    def get(self, e: Expr) -> sql.ColumnElement | None:
         """Returns the sql.ColumnElement for the given Expr, or None if Expr.to_sql() returns None."""
         try:
             return self.cache[e.id]

pixeltable/exprs/string_op.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Any
 
 import sqlalchemy as sql
 
@@ -60,7 +60,7 @@ class StringOp(Expr):
     def _id_attrs(self) -> list[tuple[str, Any]]:
         return [*super()._id_attrs(), ('operator', self.operator.value)]
 
-    def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, sql_elements: SqlElementCache) -> sql.ColumnElement | None:
         left = sql_elements.get(self._op1)
         right = sql_elements.get(self._op2)
         if left is None or right is None:

pixeltable/exprs/type_cast.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import sqlalchemy as sql
 
 from pixeltable import type_system as ts
@@ -19,7 +17,7 @@ class TypeCast(Expr):
     def __init__(self, underlying: Expr, new_type: ts.ColumnType):
         super().__init__(new_type)
         self.components: list[Expr] = [underlying]
-        self.id: Optional[int] = self._create_id()
+        self.id: int | None = self._create_id()
 
     def _equals(self, other: 'TypeCast') -> bool:
         # `TypeCast` has no properties beyond those captured by `Expr`.
@@ -29,7 +27,7 @@ class TypeCast(Expr):
     def _op1(self) -> Expr:
         return self.components[0]
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
         """
         sql_expr() is unimplemented for now, in order to sidestep potentially thorny
         questions about consistency of doing type conversions in both Python and Postgres.
@@ -40,7 +38,7 @@ class TypeCast(Expr):
         original_val = data_row[self._op1.slot_idx]
         data_row[self.slot_idx] = self.col_type.create_literal(original_val)
 
-    def as_literal(self) -> Optional[Literal]:
+    def as_literal(self) -> Literal | None:
         op1_lit = self._op1.as_literal()
         if op1_lit is None:
             return None

pixeltable/func/aggregate_function.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import abc
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Sequence, overload
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Sequence, overload
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
@@ -42,7 +42,7 @@ class AggregateFunction(Function):
     def __init__(
         self,
         agg_class: type[Aggregator],
-        type_substitutions: Optional[Sequence[dict]],
+        type_substitutions: Sequence[dict] | None,
         self_path: str,
         requires_order_by: bool,
         allows_std_agg: bool,
@@ -75,7 +75,7 @@ class AggregateFunction(Function):
         self.init_param_names = [self.init_param_names[signature_idx]]
 
     def __cls_to_signature(
-        self, cls: type[Aggregator], type_substitutions: Optional[dict] = None
+        self, cls: type[Aggregator], type_substitutions: dict | None = None
     ) -> tuple[Signature, list[str]]:
         """Inspects the Aggregator class to infer the corresponding function signature. Returns the
         inferred signature along with the list of init_param_names (for downstream error handling).
@@ -159,7 +159,7 @@ class AggregateFunction(Function):
         self.init_param_names.append(init_param_names)
         return self
 
-    def comment(self) -> Optional[str]:
+    def comment(self) -> str | None:
         return inspect.getdoc(self.agg_classes[0])
 
     def help_str(self) -> str:
@@ -173,7 +173,7 @@ class AggregateFunction(Function):
         from pixeltable import exprs
 
         # perform semantic analysis of special parameters 'order_by' and 'group_by'
-        order_by_clause: Optional[Any] = None
+        order_by_clause: Any | None = None
         if self.ORDER_BY_PARAM in kwargs:
             if self.requires_order_by:
                 raise excs.Error(
@@ -198,7 +198,7 @@ class AggregateFunction(Function):
             # don't pass the first parameter on, the Function doesn't get to see it
             args = args[1:]
 
-        group_by_clause: Optional[Any] = None
+        group_by_clause: Any | None = None
         if self.GROUP_BY_PARAM in kwargs:
             if not self.allows_window:
                 raise excs.Error(
@@ -248,7 +248,7 @@ def uda(
     requires_order_by: bool = False,
     allows_std_agg: bool = True,
     allows_window: bool = False,
-    type_substitutions: Optional[Sequence[dict]] = None,
+    type_substitutions: Sequence[dict] | None = None,
 ) -> Callable[[type[Aggregator]], AggregateFunction]: ...
 
 
@@ -302,7 +302,7 @@ def make_aggregator(
     requires_order_by: bool = False,
     allows_std_agg: bool = True,
     allows_window: bool = False,
-    type_substitutions: Optional[Sequence[dict]] = None,
+    type_substitutions: Sequence[dict] | None = None,
 ) -> AggregateFunction:
     class_path = f'{cls.__module__}.{cls.__qualname__}'
     instance = AggregateFunction(cls, type_substitutions, class_path, requires_order_by, allows_std_agg, allows_window)

pixeltable/func/callable_function.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence
+from typing import TYPE_CHECKING, Any, Callable, Sequence
 from uuid import UUID
 
 import cloudpickle  # type: ignore[import-untyped]
@@ -25,16 +25,16 @@ class CallableFunction(Function):
     """
 
     py_fns: list[Callable]
-    self_name: Optional[str]
-    batch_size: Optional[int]
+    self_name: str | None
+    batch_size: int | None
 
     def __init__(
         self,
         signatures: list[Signature],
         py_fns: list[Callable],
-        self_path: Optional[str] = None,
-        self_name: Optional[str] = None,
-        batch_size: Optional[int] = None,
+        self_path: str | None = None,
+        self_name: str | None = None,
+        batch_size: int | None = None,
         is_method: bool = False,
         is_property: bool = False,
     ):
@@ -60,7 +60,7 @@ class CallableFunction(Function):
     def is_async(self) -> bool:
         return inspect.iscoroutinefunction(self.py_fn)
 
-    def comment(self) -> Optional[str]:
+    def comment(self) -> str | None:
         return inspect.getdoc(self.py_fns[0])
 
     @property
@@ -138,7 +138,7 @@ class CallableFunction(Function):
         batched_kwargs = {k: v for k, v in kwargs.items() if k not in constant_param_names}
         return constant_kwargs, batched_kwargs
 
-    def get_batch_size(self, *args: Any, **kwargs: Any) -> Optional[int]:
+    def get_batch_size(self, *args: Any, **kwargs: Any) -> int | None:
         return self.batch_size
 
     @property
@@ -187,7 +187,7 @@ class CallableFunction(Function):
         return md, cloudpickle.dumps(self.py_fn)
 
     @classmethod
-    def from_store(cls, name: Optional[str], md: dict, binary_obj: bytes) -> Function:
+    def from_store(cls, name: str | None, md: dict, binary_obj: bytes) -> Function:
         py_fn = cloudpickle.loads(binary_obj)
         assert callable(py_fn)
         sig = Signature.from_dict(md['signature'])

pixeltable/func/expr_template_function.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Sequence
+from typing import Any, Sequence
 
 from pixeltable import exceptions as excs, exprs, type_system as ts
 
@@ -41,7 +41,7 @@ class ExprTemplateFunction(Function):
     templates: list[ExprTemplate]
     self_name: str
 
-    def __init__(self, templates: list[ExprTemplate], self_path: Optional[str] = None, name: Optional[str] = None):
+    def __init__(self, templates: list[ExprTemplate], self_path: str | None = None, name: str | None = None):
         self.templates = templates
         self.self_name = name
 
@@ -98,7 +98,7 @@ class ExprTemplateFunction(Function):
         )
         return substituted_expr.col_type
 
-    def comment(self) -> Optional[str]:
+    def comment(self) -> str | None:
         if isinstance(self.templates[0].expr, exprs.FunctionCall):
             return self.templates[0].expr.fn.comment()
         return None

pixeltable/func/function.py

@@ -5,7 +5,7 @@ import inspect
 import typing
 from abc import ABC, abstractmethod
 from copy import copy
-from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, cast
+from typing import TYPE_CHECKING, Any, Callable, Sequence, cast
 
 import sqlalchemy as sql
 from typing_extensions import Self
@@ -30,10 +30,10 @@ class Function(ABC):
     """
 
     signatures: list[Signature]
-    self_path: Optional[str]
+    self_path: str | None
     is_method: bool
     is_property: bool
-    _conditional_return_type: Optional[Callable[..., ts.ColumnType]]
+    _conditional_return_type: Callable[..., ts.ColumnType] | None
 
     # We cache the overload resolutions in self._resolutions. This ensures that each resolution is represented
     # globally by a single Python object. We do this dynamically rather than pre-constructing them in order to
@@ -43,17 +43,17 @@ class Function(ABC):
     # Translates a call to this function with the given arguments to its SQLAlchemy equivalent.
     # Overriden for specific Function instances via the to_sql() decorator. The override must accept the same
     # parameter names as the original function. Each parameter is going to be of type sql.ColumnElement.
-    _to_sql: Callable[..., Optional[sql.ColumnElement]]
+    _to_sql: Callable[..., sql.ColumnElement | None]
 
     # Returns the resource pool to use for calling this function with the given arguments.
     # Overriden for specific Function instances via the resource_pool() decorator. The override must accept a subset
     # of the parameters of the original function, with the same type.
-    _resource_pool: Callable[..., Optional[str]]
+    _resource_pool: Callable[..., str | None]
 
     def __init__(
         self,
         signatures: list[Signature],
-        self_path: Optional[str] = None,
+        self_path: str | None = None,
         is_method: bool = False,
         is_property: bool = False,
     ):
@@ -105,7 +105,7 @@ class Function(ABC):
     @abstractmethod
     def is_async(self) -> bool: ...
 
-    def comment(self) -> Optional[str]:
+    def comment(self) -> str | None:
         return None
 
     def help_str(self) -> str:
@@ -175,7 +175,7 @@ class Function(ABC):
 
     def _bind_to_matching_signature(self, args: Sequence[Any], kwargs: dict[str, Any]) -> tuple[Self, dict[str, Any]]:
         result: int = -1
-        bound_args: Optional[dict[str, Any]] = None
+        bound_args: dict[str, Any] | None = None
         assert len(self.signatures) > 0
         if len(self.signatures) == 1:
             # Only one signature: call _bind_to_signature() and surface any errors directly
@@ -206,7 +206,7 @@ class Function(ABC):
         self._resolved_fns[signature_idx].validate_call(normalized_args)
         return normalized_args
 
-    def validate_call(self, bound_args: dict[str, Optional['exprs.Expr']]) -> None:
+    def validate_call(self, bound_args: dict[str, 'exprs.Expr' | None]) -> None:
         """Override this to do custom validation of the arguments"""
         assert not self.is_polymorphic
         self.signature.validate_args(bound_args, context=f'in function {self.name!r}')
@@ -252,9 +252,7 @@ class Function(ABC):
 
         return return_type
 
-    def _assemble_callable_args(
-        self, callable: Callable, bound_args: dict[str, 'exprs.Expr']
-    ) -> Optional[dict[str, Any]]:
+    def _assemble_callable_args(self, callable: Callable, bound_args: dict[str, 'exprs.Expr']) -> dict[str, Any] | None:
         """
         Return the kwargs to pass to callable, given bound_args passed to this function.
 
@@ -286,7 +284,7 @@ class Function(ABC):
                 return None
             assert isinstance(arg, exprs.Expr)
 
-            expects_expr: Optional[type[exprs.Expr]] = None
+            expects_expr: type[exprs.Expr] | None = None
             type_hint = callable_type_hints.get(param.name)
             if typing.get_origin(type_hint) is not None:
                 type_hint = typing.get_origin(type_hint)  # Remove type subscript if one exists
@@ -410,12 +408,12 @@ class Function(ABC):
         """Execute the function with the given arguments and return the result."""
         raise NotImplementedError()
 
-    def to_sql(self, fn: Callable[..., Optional[sql.ColumnElement]]) -> Callable[..., Optional[sql.ColumnElement]]:
+    def to_sql(self, fn: Callable[..., sql.ColumnElement | None]) -> Callable[..., sql.ColumnElement | None]:
         """Instance decorator for specifying the SQL translation of this function"""
         self._to_sql = fn
         return fn
 
-    def __default_to_sql(self, *args: Any, **kwargs: Any) -> Optional[sql.ColumnElement]:
+    def __default_to_sql(self, *args: Any, **kwargs: Any) -> sql.ColumnElement | None:
         """The default implementation of SQL translation, which provides no translation"""
         return None
 
@@ -425,7 +423,7 @@ class Function(ABC):
         self._resource_pool = fn
         return fn
 
-    def __default_resource_pool(self) -> Optional[str]:
+    def __default_resource_pool(self) -> str | None:
         return None
 
     def __eq__(self, other: object) -> bool:
@@ -495,7 +493,7 @@ class Function(ABC):
         raise NotImplementedError()
 
     @classmethod
-    def from_store(cls, name: Optional[str], md: dict, binary_obj: bytes) -> Function:
+    def from_store(cls, name: str | None, md: dict, binary_obj: bytes) -> Function:
         """
         Create a Function instance from the serialized representation returned by to_store()
         """

pixeltable/func/function_registry.py

@@ -4,7 +4,6 @@ import dataclasses
 import importlib
 import logging
 import sys
-from typing import Optional
 from uuid import UUID
 
 import sqlalchemy as sql
@@ -23,7 +22,7 @@ class FunctionRegistry:
     Function are loaded from the store on demand.
     """
 
-    _instance: Optional[FunctionRegistry] = None
+    _instance: FunctionRegistry | None = None
 
     @classmethod
     def get(cls) -> FunctionRegistry:
@@ -94,7 +93,7 @@ class FunctionRegistry:
         #         stored_fn_md.append(md)
         return list(self.module_fns.values())
 
-    # def get_function(self, *, id: Optional[UUID] = None, fqn: Optional[str] = None) -> Function:
+    # def get_function(self, *, id: UUID | None = None, fqn: str | None = None) -> Function:
     #     assert (id is not None) != (fqn is not None)
     #     if id is not None:
     #         if id not in self.stored_fns_by_id:
@@ -143,7 +142,7 @@ class FunctionRegistry:
             return list(self.type_methods[base_type].values())
         return []
 
-    def lookup_type_method(self, base_type: ts.ColumnType.Type, name: str) -> Optional[Function]:
+    def lookup_type_method(self, base_type: ts.ColumnType.Type, name: str) -> Function | None:
         """
         Look up a method (or property) by name for a given base type. If no such method is registered, return None.
         """
@@ -151,8 +150,8 @@ class FunctionRegistry:
             return self.type_methods[base_type][name]
         return None
 
-    # def create_function(self, md: schema.FunctionMd, binary_obj: bytes, dir_id: Optional[UUID] = None) -> UUID:
-    def create_stored_function(self, pxt_fn: Function, dir_id: Optional[UUID] = None) -> UUID:
+    # def create_function(self, md: schema.FunctionMd, binary_obj: bytes, dir_id: UUID | None = None) -> UUID:
+    def create_stored_function(self, pxt_fn: Function, dir_id: UUID | None = None) -> UUID:
         fn_md, binary_obj = pxt_fn.to_store()
         md = schema.FunctionMd(name=pxt_fn.name, md=fn_md, py_version=sys.version, class_name=pxt_fn.__class__.__name__)
         with env.Env.get().engine.begin() as conn:
@@ -184,7 +183,7 @@ class FunctionRegistry:
             self.stored_fns_by_id[id] = instance
             return instance
 
-    # def create_function(self, fn: Function, dir_id: Optional[UUID] = None, name: Optional[str] = None) -> None:
+    # def create_function(self, fn: Function, dir_id: UUID | None = None, name: str | None = None) -> None:
     #     with env.Env.get().engine.begin() as conn:
     #         _logger.debug(f'Pickling function {name}')
     #         eval_fn_str = cloudpickle.dumps(fn.eval_fn) if fn.eval_fn is not None else None

pixeltable/func/globals.py

@@ -1,14 +1,13 @@
 import importlib
 import inspect
 from types import ModuleType
-from typing import Optional
 
 import pixeltable.exceptions as excs
 
 
-def resolve_symbol(symbol_path: str) -> Optional[object]:
+def resolve_symbol(symbol_path: str) -> object | None:
     path_elems = symbol_path.split('.')
-    module: Optional[ModuleType] = None
+    module: ModuleType | None = None
     i = len(path_elems) - 1
     while i > 0 and module is None:
         try:

pixeltable/func/mcp.py

@@ -1,6 +1,6 @@
 import asyncio
 import inspect
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 import pixeltable as pxt
 from pixeltable import exceptions as excs, type_system as ts
@@ -18,7 +18,7 @@ async def mcp_udfs_async(url: str) -> list['pxt.func.Function']:
     import mcp
     from mcp.client.streamable_http import streamablehttp_client
 
-    list_tools_result: Optional[mcp.types.ListToolsResult] = None
+    list_tools_result: mcp.types.ListToolsResult | None = None
     async with (
         streamablehttp_client(url) as (read_stream, write_stream, _),
         mcp.ClientSession(read_stream, write_stream) as session,

pixeltable/func/query_template_function.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import inspect
 from functools import reduce
-from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional, overload
+from typing import TYPE_CHECKING, Any, Callable, Iterable, overload
 
 from pixeltable import catalog, exceptions as excs, exprs, func, type_system as ts
 
@@ -16,13 +16,13 @@ if TYPE_CHECKING:
 class QueryTemplateFunction(Function):
     """A parameterized query/DataFrame from which an executable DataFrame is created with a function call."""
 
-    template_df: Optional['DataFrame']
-    self_name: Optional[str]
-    _comment: Optional[str]
+    template_df: 'DataFrame' | None
+    self_name: str | None
+    _comment: str | None
 
     @classmethod
     def create(
-        cls, template_callable: Callable, param_types: Optional[list[ts.ColumnType]], path: str, name: str
+        cls, template_callable: Callable, param_types: list[ts.ColumnType] | None, path: str, name: str
     ) -> QueryTemplateFunction:
         # we need to construct a template df and a signature
         py_sig = inspect.signature(template_callable)
@@ -40,11 +40,11 @@ class QueryTemplateFunction(Function):
 
     def __init__(
         self,
-        template_df: Optional['DataFrame'],
+        template_df: 'DataFrame' | None,
         sig: Signature,
-        path: Optional[str] = None,
-        name: Optional[str] = None,
-        comment: Optional[str] = None,
+        path: str | None = None,
+        name: str | None = None,
+        comment: str | None = None,
     ):
         assert sig is not None
         super().__init__([sig], self_path=path)
@@ -82,7 +82,7 @@ class QueryTemplateFunction(Function):
     def name(self) -> str:
         return self.self_name
 
-    def comment(self) -> Optional[str]:
+    def comment(self) -> str | None:
         return self._comment
 
     def _as_dict(self) -> dict:
@@ -100,11 +100,11 @@ def query(py_fn: Callable) -> QueryTemplateFunction: ...
 
 
 @overload
-def query(*, param_types: Optional[list[ts.ColumnType]] = None) -> Callable[[Callable], QueryTemplateFunction]: ...
+def query(*, param_types: list[ts.ColumnType] | None = None) -> Callable[[Callable], QueryTemplateFunction]: ...
 
 
 def query(*args: Any, **kwargs: Any) -> Any:
-    def make_query_template(py_fn: Callable, param_types: Optional[list[ts.ColumnType]]) -> QueryTemplateFunction:
+    def make_query_template(py_fn: Callable, param_types: list[ts.ColumnType] | None) -> QueryTemplateFunction:
         if py_fn.__module__ != '__main__' and py_fn.__name__.isidentifier():
             # this is a named function in a module
             function_path = f'{py_fn.__module__}.{py_fn.__qualname__}'
@@ -127,10 +127,10 @@ def query(*args: Any, **kwargs: Any) -> Any:
 
 def retrieval_udf(
     table: catalog.Table,
-    name: Optional[str] = None,
-    description: Optional[str] = None,
-    parameters: Optional[Iterable[str | exprs.ColumnRef]] = None,
-    limit: Optional[int] = 10,
+    name: str | None = None,
+    description: str | None = None,
+    parameters: Iterable[str | exprs.ColumnRef] | None = None,
+    limit: int | None = 10,
 ) -> func.QueryTemplateFunction:
     """
     Constructs a retrieval UDF for the given table. The retrieval UDF is a UDF whose parameters are