pixeltable 0.2.15__py3-none-any.whl → 0.2.16__py3-none-any.whl

pixeltable/exprs/expr.py

@@ -7,7 +7,7 @@ import inspect
 import json
 import sys
 import typing
-from typing import Union, Optional, List, Callable, Any, Dict, Tuple, Set, Generator, Type
+from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Type, TypeVar, Union, overload
 from uuid import UUID
 
 import sqlalchemy as sql
@@ -17,8 +17,9 @@ import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
 import pixeltable.func as func
 import pixeltable.type_system as ts
+
 from .data_row import DataRow
-from .globals import ComparisonOperator, LogicalOperator, LiteralPythonTypes, ArithmeticOperator
+from .globals import ArithmeticOperator, ComparisonOperator, LiteralPythonTypes, LogicalOperator
 
 
 class ExprScope:
@@ -191,8 +192,8 @@ class Expr(abc.ABC):
         Recursively replace ColRefs to unstored computed columns with their value exprs.
         Also replaces references to stored computed columns in resolve_cols.
         """
-        from .expr_set import ExprSet
         from .column_ref import ColumnRef
+        from .expr_set import ExprSet
         if resolve_cols is None:
             resolve_cols = set()
         result = self
@@ -240,28 +241,64 @@ class Expr(abc.ABC):
         return str(self)
 
     @classmethod
-    def print_list(cls, expr_list: List[Expr]) -> str:
+    def print_list(cls, expr_list: list[Any]) -> str:
         if len(expr_list) == 1:
             return str(expr_list[0])
-        return f'({", ".join([str(e) for e in expr_list])})'
+        return f'({", ".join(str(e) for e in expr_list)})'
+
+    # `subexprs` has two forms: one that takes an explicit subclass of `Expr` as an argument and returns only
+    # instances of that subclass; and another that returns all subexpressions that match the given filter.
+    # In order for type checking to behave correctly on both forms, we provide two overloaded signatures.
+
+    T = TypeVar('T', bound='Expr')
+
+    @overload
+    def subexprs(
+        self, *, filter: Optional[Callable[[Expr], bool]] = None, traverse_matches: bool = True
+    ) -> Iterator[Expr]: ...
 
+    @overload
     def subexprs(
-            self, expr_class: Optional[Type[Expr]] = None, filter: Optional[Callable[[Expr], bool]] = None,
-            traverse_matches: bool = True
-    ) -> Generator[Expr, None, None]:
+        self, expr_class: type[T], filter: Optional[Callable[[Expr], bool]] = None,
+        traverse_matches: bool = True
+    ) -> Iterator[T]: ...
+
+    def subexprs(
+        self, expr_class: Optional[type[T]] = None, filter: Optional[Callable[[Expr], bool]] = None,
+        traverse_matches: bool = True
+    ) -> Iterator[T]:
         """
         Iterate over all subexprs, including self.
         """
-        assert expr_class is None or filter is None  # at most one of them
-        if expr_class is not None:
-            filter = lambda e: isinstance(e, expr_class)
         is_match = filter is None or filter(self)
+        if expr_class is not None:
+            is_match = is_match and isinstance(self, expr_class)
         if not is_match or traverse_matches:
             for c in self.components:
-                yield from c.subexprs(filter=filter, traverse_matches=traverse_matches)
+                yield from c.subexprs(expr_class=expr_class, filter=filter, traverse_matches=traverse_matches)
         if is_match:
             yield self
 
+    @overload
+    def list_subexprs(
+        expr_list: list[Expr], *, filter: Optional[Callable[[Expr], bool]] = None, traverse_matches: bool = True
+    ) -> Iterator[Expr]: ...
+
+    @overload
+    def list_subexprs(
+        expr_list: list[Expr], expr_class: type[T], filter: Optional[Callable[[Expr], bool]] = None,
+        traverse_matches: bool = True
+    ) -> Iterator[T]: ...
+
+    @classmethod
+    def list_subexprs(
+        cls, expr_list: list[Expr], expr_class: Optional[type[T]] = None,
+        filter: Optional[Callable[[Expr], bool]] = None, traverse_matches: bool = True
+    ) -> Iterator[T]:
+        """Produce subexprs for all exprs in list. Can contain duplicates."""
+        for e in expr_list:
+            yield from e.subexprs(expr_class=expr_class, filter=filter, traverse_matches=traverse_matches)
+
     def _contains(self, cls: Optional[Type[Expr]] = None, filter: Optional[Callable[[Expr], bool]] = None) -> bool:
         """
         Returns True if any subexpr is an instance of cls.
@@ -275,15 +312,6 @@ class Expr(abc.ABC):
         except StopIteration:
             return False
 
-    @classmethod
-    def list_subexprs(
-            cls, expr_list: List[Expr], expr_class: Optional[Type[Expr]] = None,
-            filter: Optional[Callable[[Expr], bool]] = None, traverse_matches: bool = True
-    ) -> Generator[Expr, None, None]:
-        """Produce subexprs for all exprs in list. Can contain duplicates."""
-        for e in expr_list:
-            yield from e.subexprs(expr_class=expr_class, filter=filter, traverse_matches=traverse_matches)
-
     def tbl_ids(self) -> Set[UUID]:
         """Returns table ids referenced by this expr."""
         from .column_ref import ColumnRef
@@ -334,10 +362,10 @@ class Expr(abc.ABC):
         return None
 
     @abc.abstractmethod
-    def sql_expr(self) -> Optional[sql.ClauseElement]:
+    def sql_expr(self) -> Optional[sql.ColumnElement]:
         """
         If this expr can be materialized directly in SQL:
-        - returns a ClauseElement
+        - returns a ColumnElement
         - eval() will not be called (exception: Literal)
         Otherwise
         - returns None

pixeltable/exprs/function_call.py

@@ -204,7 +204,14 @@ class FunctionCall(Expr):
                 # Check that the argument is consistent with the expected parameter type, with the allowance that
                 # non-nullable parameters can still accept nullable arguments (since function calls with Nones
                 # assigned to non-nullable parameters will always return None)
-                if not param.col_type.is_supertype_of(arg.col_type, ignore_nullable=True):
+                if not (
+                    param.col_type.is_supertype_of(arg.col_type, ignore_nullable=True)
+                    # TODO: this is a hack to allow JSON columns to be passed to functions that accept scalar
+                    # types. It's necessary to avoid littering notebooks with `apply(str)` calls or equivalent.
+                    # (Previously, this wasn't necessary because `is_supertype_of()` was improperly implemented.)
+                    # We need to think through the right way to handle this scenario.
+                    or (arg.col_type.is_json_type() and param.col_type.is_scalar_type())
+                ):
                     raise excs.Error(
                         f'Parameter {param_name}: argument type {arg.col_type} does not match parameter type '
                         f'{param.col_type}')

pixeltable/exprs/inline_array.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import copy
-from typing import Optional, List, Any, Dict, Tuple
+from typing import Iterable, Optional, List, Any, Dict, Tuple
 
 import numpy as np
 import sqlalchemy as sql
@@ -24,7 +24,7 @@ class InlineArray(Expr):
 
     elements: List[Tuple[Optional[int], Any]]
 
-    def __init__(self, elements: Tuple, force_json: bool = False):
+    def __init__(self, elements: Iterable, force_json: bool = False):
         # we need to call this in order to populate self.components
         super().__init__(ts.ArrayType((len(elements),), ts.IntType()))
 

pixeltable/exprs/json_path.py

@@ -1,22 +1,29 @@
 from __future__ import annotations
-from typing import Optional, List, Any, Dict, Tuple, Union
+
+from typing import Any, Optional, Union
 
 import jmespath
 import sqlalchemy as sql
 
-from .globals import print_slice
-from .expr import Expr
-from .json_mapper import JsonMapper
-from .data_row import DataRow
-from .row_builder import RowBuilder
 import pixeltable
-import pixeltable.exceptions as excs
 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
+
 
 class JsonPath(Expr):
-    def __init__(self, anchor: Optional['pixeltable.exprs.ColumnRef'], path_elements: Optional[List[str]] = None, scope_idx: int = 0):
+    def __init__(
+        self,
+        anchor: Optional['pixeltable.exprs.ColumnRef'],
+        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().
         scope_idx: for relative paths, index of referenced JsonMapper
@@ -27,7 +34,7 @@ class JsonPath(Expr):
         super().__init__(ts.JsonType())
         if anchor is not None:
             self.components = [anchor]
-        self.path_elements: List[Union[str, int]] = path_elements
+        self.path_elements: list[Union[str, int, slice]] = path_elements
         self.compiled_path = jmespath.compile(self._json_path()) if len(path_elements) > 0 else None
         self.scope_idx = scope_idx
         # NOTE: the _create_id() result will change if set_anchor() gets called;
@@ -39,16 +46,26 @@ class JsonPath(Expr):
         return (f'{str(self._anchor) if self._anchor is not None else "R"}'
             f'{"." if isinstance(self.path_elements[0], str) else ""}{self._json_path()}')
 
-    def _as_dict(self) -> Dict:
-        return {'path_elements': self.path_elements, 'scope_idx': self.scope_idx, **super()._as_dict()}
+    def _as_dict(self) -> dict:
+        path_elements = [
+            [el.start, el.stop, el.step] if isinstance(el, slice)
+            else el
+            for el in self.path_elements
+        ]
+        return {'path_elements': path_elements, 'scope_idx': self.scope_idx, **super()._as_dict()}
 
     @classmethod
-    def _from_dict(cls, d: Dict, components: List[Expr]) -> Expr:
+    def _from_dict(cls, d: dict, components: list[Expr]) -> Expr:
         assert 'path_elements' in d
         assert 'scope_idx' in d
         assert len(components) <= 1
         anchor = components[0] if len(components) == 1 else None
-        return cls(anchor, d['path_elements'], d['scope_idx'])
+        path_elements = [
+            slice(el[0], el[1], el[2]) if isinstance(el, list)
+            else el
+            for el in d['path_elements']
+        ]
+        return cls(anchor, path_elements, d['scope_idx'])
 
     @property
     def _anchor(self) -> Optional[Expr]:
@@ -85,8 +102,7 @@ class JsonPath(Expr):
         if isinstance(index, str):
             if index != '*':
                 raise excs.Error(f'Invalid json list index: {index}')
-        else:
-            if not isinstance(index, slice) and not isinstance(index, int):
+        elif not isinstance(index, (int, slice)):
                 raise excs.Error(f'Invalid json list index: {index}')
         return JsonPath(self._anchor, self.path_elements + [index])
 
@@ -99,7 +115,7 @@ class JsonPath(Expr):
     def default_column_name(self) -> Optional[str]:
         anchor_name = self._anchor.default_column_name() if self._anchor is not None else ''
         ret_name = f'{anchor_name}.{self._json_path()}'
-        
+
         def cleanup_char(s : str) -> str:
             if s == '.':
                 return '_'
@@ -109,19 +125,19 @@ class JsonPath(Expr):
                 return s
             else:
                 return ''
-            
+
         clean_name = ''.join(map(cleanup_char, ret_name))
         clean_name = clean_name.lstrip('_') # remove leading underscore
         if clean_name == '':
             clean_name = None
-        
+
         assert clean_name is None or catalog.is_valid_identifier(clean_name)
         return clean_name
 
     def _equals(self, other: JsonPath) -> bool:
         return self.path_elements == other.path_elements
 
-    def _id_attrs(self) -> List[Tuple[str, Any]]:
+    def _id_attrs(self) -> list[tuple[str, Any]]:
         return super()._id_attrs() + [('path_elements', self.path_elements)]
 
     def sql_expr(self) -> Optional[sql.ClauseElement]:
@@ -137,7 +153,7 @@ class JsonPath(Expr):
 
     def _json_path(self) -> str:
         assert len(self.path_elements) > 0
-        result: List[str] = []
+        result: list[str] = []
         for element in self.path_elements:
             if element == '*':
                 result.append('[*]')

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, Tuple, Set
+from typing import Optional, List, Any, Dict, Sequence, Tuple, Set
 
 import sqlalchemy as sql
 
@@ -57,7 +57,7 @@ class RowBuilder:
         target_exprs: List[Expr]  # exprs corresponding to target_slot_idxs
 
     def __init__(
-            self, output_exprs: List[Expr], columns: List[catalog.Column], input_exprs: List[Expr]
+            self, output_exprs: Sequence[Expr], columns: Sequence[catalog.Column], input_exprs: Sequence[Expr]
     ):
         """
         Args:
@@ -227,10 +227,10 @@ 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[Expr], remove_duplicates: bool = True) -> None:
+    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[i] = set()  # slot idxs within expr_list
+        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:

pixeltable/exprs/rowid_ref.py

@@ -56,7 +56,7 @@ class RowidRef(Expr):
         # check if this is the pos column of a component view
         tbl = self.tbl if self.tbl is not None else catalog.Catalog.get().tbl_versions[(self.tbl_id, None)]
         if tbl.is_component_view() and self.rowid_component_idx == tbl.store_tbl.pos_col_idx:
-            return catalog.globals.POS_COLUMN_NAME
+            return catalog.globals._POS_COLUMN_NAME
         return ''
 
     def set_tbl(self, tbl: catalog.TableVersionPath) -> None:

pixeltable/functions/__init__.py

@@ -1,8 +1,7 @@
-from . import fireworks, huggingface, image, openai, string, together, video, timestamp, json, vision
+from . import audio, fireworks, huggingface, image, json, openai, string, timestamp, together, video, vision
 from .globals import *
 from pixeltable.utils.code import local_public_names
 
-
 __all__ = local_public_names(__name__, exclude=['globals']) + local_public_names(globals.__name__)
 
 

pixeltable/functions/audio.py

@@ -0,0 +1,32 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for `AudioType`.
+
+Example:
+```python
+import pixeltable as pxt
+import pixeltable.functions as pxtf
+
+t = pxt.get_table(...)
+t.select(pxtf.audio.get_metadata()).collect()
+```
+"""
+
+import pixeltable.func as func
+import pixeltable.type_system as ts
+from pixeltable.utils.code import local_public_names
+
+
+@func.udf(return_type=ts.JsonType(nullable=False), param_types=[ts.AudioType(nullable=False)], is_method=True)
+def get_metadata(audio: str) -> dict:
+    """
+    Gets various metadata associated with an audio file and returns it as a dictionary.
+    """
+    import pixeltable.functions as pxtf
+    return pxtf.video._get_metadata(audio)
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/huggingface.py

@@ -50,7 +50,7 @@ def sentence_transformer(
         >>> tbl['result'] = sentence_transformer(tbl.sentence, model_id='all-mpnet-base-v2')
     """
     env.Env.get().require_package('sentence_transformers')
-    from sentence_transformers import SentenceTransformer
+    from sentence_transformers import SentenceTransformer  # type: ignore
 
     model = _lookup_model(model_id, SentenceTransformer)
 
@@ -154,7 +154,7 @@ def clip_text(text: Batch[str], *, model_id: str) -> Batch[np.ndarray]:
     env.Env.get().require_package('transformers')
     device = resolve_torch_device('auto')
     import torch
-    from transformers import CLIPModel, CLIPProcessor
+    from transformers import CLIPModel, CLIPProcessor  # type: ignore
 
     model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
     processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
@@ -317,9 +317,9 @@ def _lookup_model(model_id: str, create: Callable[[str], T], device: Optional[st
     key = (model_id, create, device)  # For safety, include the `create` callable in the cache key
     if key not in _model_cache:
         model = create(model_id)
-        if device is not None:
-            model.to(device)
         if isinstance(model, nn.Module):
+            if device is not None:
+                model.to(device)
             model.eval()
         _model_cache[key] = model
     return _model_cache[key]

pixeltable/functions/image.py

@@ -144,7 +144,7 @@ def resize(self: PIL.Image.Image, size: tuple[int, int]) -> PIL.Image.Image:
     Equivalent to
     [`PIL.Image.Image.resize()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.resize)
     """
-    return self.resize(size)
+    return self.resize(tuple(size))
 
 
 @resize.conditional_return_type

pixeltable/functions/video.py

@@ -133,7 +133,11 @@ def get_metadata(video: str) -> dict:
     """
     Gets various metadata associated with a video file and returns it as a dictionary.
     """
-    with av.open(video) as container:
+    return _get_metadata(video)
+
+
+def _get_metadata(path: str) -> dict:
+    with av.open(path) as container:
         assert isinstance(container, av.container.InputContainer)
         streams_info = [__get_stream_metadata(stream) for stream in container.streams]
         result = {

pixeltable/functions/vision.py

@@ -274,7 +274,6 @@ def draw_bounding_boxes(
         boxes: list[list[int]],
         labels: Optional[list[Any]] = None,
         color: Optional[str] = None,
-        label_colors: Optional[dict[Union[str, int], str]] = None,
         box_colors: Optional[list[str]] = None,
         fill: bool  = False,
         width: int = 1,
@@ -297,7 +296,6 @@ def draw_bounding_boxes(
         boxes: List of bounding boxes, each represented as [xmin, ymin, xmax, ymax].
         labels: List of labels for each bounding box.
         color: Single color to be used for all bounding boxes and labels.
-        label_colors: Dictionary mapping labels to colors.
         box_colors: List of colors, one per bounding box.
         fill: Whether to fill the bounding boxes with color.
         width: Width of the bounding box borders.
@@ -310,9 +308,9 @@ def draw_bounding_boxes(
     Returns:
         The image with bounding boxes drawn on it.
     """
-    color_params = sum([color is not None, label_colors is not None, box_colors is not None])
+    color_params = sum([color is not None, box_colors is not None])
     if color_params > 1:
-        raise ValueError("Only one of 'color', 'label_colors', or 'box_colors' can be set")
+        raise ValueError("Only one of 'color' or 'box_colors' can be set")
 
     # ensure the number of labels matches the number of boxes
     num_boxes = len(boxes)
@@ -328,8 +326,6 @@ def draw_bounding_boxes(
     else:
         if color is not None:
             box_colors = [color] * num_boxes
-        elif label_colors is not None:
-            box_colors = [label_colors.get(label, DEFAULT_COLOR) for label in labels]
         else:
             label_colors = _create_label_colors(labels)
             box_colors = [label_colors[label] for label in labels]