pixeltable 0.2.12__py3-none-any.whl → 0.2.14__py3-none-any.whl

pixeltable/exprs/in_predicate.py

@@ -5,20 +5,20 @@ from typing import Optional, List, Any, Dict, Tuple, Iterable
 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 .predicate import Predicate
 from .row_builder import RowBuilder
 
 
-class InPredicate(Predicate):
+class InPredicate(Expr):
     """Predicate corresponding to the SQL IN operator."""
 
     def __init__(self, lhs: Expr, value_set_literal: Optional[Iterable] = None, value_set_expr: Optional[Expr] = None):
         assert (value_set_literal is None) != (value_set_expr is None)
         if not lhs.col_type.is_scalar_type():
             raise excs.Error(f'isin(): only supported for scalar types, not {lhs.col_type}')
-        super().__init__()
+        super().__init__(ts.BoolType())
 
         self.value_list: Optional[list] = None  # only contains values of the correct type
         if value_set_expr is not None:

pixeltable/exprs/inline_array.py

@@ -21,6 +21,9 @@ class InlineArray(Expr):
     is `True`, it will always be cast as a `JsonType`. If `force_json` is `False`, it will be cast as an
     `ArrayType` if it is a homogenous array of scalars or arrays, or a `JsonType` otherwise.
     """
+
+    elements: List[Tuple[Optional[int], Any]]
+
     def __init__(self, elements: Tuple, force_json: bool = False):
         # we need to call this in order to populate self.components
         super().__init__(ts.ArrayType((len(elements),), ts.IntType()))
@@ -28,7 +31,7 @@ class InlineArray(Expr):
         # elements contains
         # - for Expr elements: (index into components, None)
         # - for non-Expr elements: (None, value)
-        self.elements: List[Tuple[Optional[int], Any]] = []
+        self.elements = []
         for el in elements:
             el = copy.deepcopy(el)
             if isinstance(el, list):
@@ -43,14 +46,16 @@ class InlineArray(Expr):
             else:
                 self.elements.append((None, el))
 
-        inferred_element_type = ts.InvalidType()
-        for idx, val in self.elements:
-            if idx is not None:
-                inferred_element_type = ts.ColumnType.supertype(inferred_element_type, self.components[idx].col_type)
-            else:
-                inferred_element_type = ts.ColumnType.supertype(inferred_element_type, ts.ColumnType.infer_literal_type(val))
-            if inferred_element_type is None:
-                break
+        inferred_element_type: Optional[ts.ColumnType] = ts.InvalidType()
+        if not force_json:
+            # try to infer the element type
+            for idx, val in self.elements:
+                if idx is not None:
+                    inferred_element_type = ts.ColumnType.supertype(inferred_element_type, self.components[idx].col_type)
+                else:
+                    inferred_element_type = ts.ColumnType.supertype(inferred_element_type, ts.ColumnType.infer_literal_type(val))
+                if inferred_element_type is None:
+                    break
 
         if force_json or inferred_element_type is None:
             # JSON conversion is forced, or there is no common supertype
@@ -93,7 +98,7 @@ class InlineArray(Expr):
             data_row[self.slot_idx] = result
 
     def _as_dict(self) -> Dict:
-        return {'elements': self.elements, **super()._as_dict()}
+        return {'elements': self.elements, 'is_json': self.col_type.is_json_type(), **super()._as_dict()}
 
     @classmethod
     def _from_dict(cls, d: Dict, components: List[Expr]) -> Expr:
@@ -106,4 +111,6 @@ class InlineArray(Expr):
                 arg.append(components[idx])
             else:
                 arg.append(val)
-        return cls(tuple(arg))
+        # in order to avoid a schema version change, we'll interpret the absence of 'is_json' to indicate an ArrayType
+        is_json = d.get('is_json', False)
+        return cls(tuple(arg), force_json=is_json)

pixeltable/exprs/is_null.py

@@ -1,18 +1,18 @@
 from __future__ import annotations
+
 from typing import Optional, List, Dict
 
 import sqlalchemy as sql
 
-from .predicate import Predicate
-from .expr import Expr
+import pixeltable.type_system as ts
 from .data_row import DataRow
+from .expr import Expr
 from .row_builder import RowBuilder
-import pixeltable.catalog as catalog
 
 
-class IsNull(Predicate):
+class IsNull(Expr):
     def __init__(self, e: Expr):
-        super().__init__()
+        super().__init__(ts.BoolType())
         self.components = [e]
         self.id = self._create_id()
 

pixeltable/exprs/method_ref.py

@@ -0,0 +1,63 @@
+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 .data_row import DataRow
+from .row_builder import RowBuilder
+
+
+class MethodRef(Expr):
+    """
+    A method reference. This represents a `Function` instance with its
+    first parameter bound to a base expression.
+
+    When a `MethodRef` is called, it returns a `FunctionCall` with the base expression as the first argument.
+    The effective arity of a `MethodRef` is one less than the arity of the underlying `Function`.
+    """
+    # 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):
+        super().__init__(ts.InvalidType())  # The `MethodRef` is untyped until it is called.
+        self.base_expr = base_expr
+        self.method_name = method_name
+        self.fn = FunctionRegistry.get().lookup_type_method(base_expr.col_type.type_enum, method_name)
+        if self.fn is None:
+            # This has to be an `AttributeError`, or tab-completion won't work properly in ipython.
+            raise AttributeError(f'Unknown method (of type {base_expr.col_type}): {method_name}')
+        self.components = [base_expr]
+        self.id = self._create_id()
+
+    def _as_dict(self) -> dict:
+        return {'method_name': self.method_name, **super()._as_dict()}
+
+    @classmethod
+    def _from_dict(cls, d: dict, components: list[Expr]) -> Expr:
+        assert 'method_name' in d
+        assert len(components) == 1
+        return cls(d['method_name'], components[0])
+
+    def __call__(self, *args, **kwargs) -> FunctionCall:
+        result = self.fn(*[self.base_expr, *args], **kwargs)
+        assert isinstance(result, FunctionCall)
+        result.is_method_call = True
+        return result
+
+    def _equals(self, other: 'MethodRef') -> bool:
+        return self.base_expr == other.base_expr and self.method_name == other.method_name
+
+    def _id_attrs(self) -> list[tuple[str, Any]]:
+        return super()._id_attrs() + [('method_name', self.method_name)]
+
+    def sql_expr(self) -> Optional[sql.ClauseElement]:
+        return None
+
+    def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
+        assert False, 'MethodRef cannot be evaluated directly'
+
+    def __str__(self) -> str:
+        return f'{self.base_expr}.{self.method_name}'

pixeltable/ext/__init__.py

@@ -3,3 +3,12 @@ Extended integrations for Pixeltable. This package contains experimental or demo
 are not intended for production use. Long-term support cannot be guaranteed, usually because the features
 have dependencies whose future support is unclear.
 """
+
+from pixeltable.utils.code import local_public_names
+from . import functions
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/ext/functions/__init__.py

@@ -0,0 +1,8 @@
+from pixeltable.utils.code import local_public_names
+from . import whisperx, yolox
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/ext/functions/whisperx.py

@@ -1,8 +1,9 @@
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
-import torch
-import whisperx
-from whisperx.asr import FasterWhisperPipeline
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    from whisperx.asr import FasterWhisperPipeline
 
 import pixeltable as pxt
 
@@ -11,6 +12,36 @@ import pixeltable as pxt
 def transcribe(
     audio: str, *, model: str, compute_type: Optional[str] = None, language: Optional[str] = None, chunk_size: int = 30
 ) -> dict:
+    """
+    Transcribe an audio file using WhisperX.
+
+    This UDF runs a transcription model _locally_ using the WhisperX library,
+    equivalent to the WhisperX `transcribe` function, as described in the
+    [WhisperX library documentation](https://github.com/m-bain/whisperX).
+
+    __Requirements:__
+
+    - `pip install whisperx`
+
+    Args:
+        audio: The audio file to transcribe.
+        model: The name of the model to use for transcription.
+
+    See the [WhisperX library documentation](https://github.com/m-bain/whisperX) for details
+    on the remaining parameters.
+
+    Returns:
+        A dictionary containing the transcription and various other metadata.
+
+    Examples:
+        Add a computed column that applies the model `tiny.en` to an existing Pixeltable column `tbl.audio`
+        of the table `tbl`:
+
+        >>> tbl['result'] = transcribe(tbl.audio, model='tiny.en')
+    """
+    import torch
+    import whisperx
+
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     compute_type = compute_type or ('float16' if device == 'cuda' else 'int8')
     model = _lookup_model(model, device, compute_type)
@@ -19,7 +50,9 @@ def transcribe(
     return result
 
 
-def _lookup_model(model_id: str, device: str, compute_type: str) -> FasterWhisperPipeline:
+def _lookup_model(model_id: str, device: str, compute_type: str) -> 'FasterWhisperPipeline':
+    import whisperx
+
     key = (model_id, device, compute_type)
     if key not in _model_cache:
         model = whisperx.load_model(model_id, device, compute_type=compute_type)
@@ -28,3 +61,10 @@ def _lookup_model(model_id: str, device: str, compute_type: str) -> FasterWhispe
 
 
 _model_cache = {}
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/ext/functions/yolox.py

@@ -1,20 +1,21 @@
 import logging
 from pathlib import Path
-from typing import Iterable, Iterator
+from typing import Iterable, Iterator, TYPE_CHECKING
 from urllib.request import urlretrieve
 
 import PIL.Image
 import numpy as np
-import torch
-from yolox.data import ValTransform
-from yolox.exp import get_exp, Exp
-from yolox.models import YOLOX
-from yolox.utils import postprocess
 
 import pixeltable as pxt
 from pixeltable import env
 from pixeltable.func import Batch
 from pixeltable.functions.util import normalize_image_mode
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import torch
+    from yolox.exp import Exp
+    from yolox.models import YOLOX
 
 _logger = logging.getLogger('pixeltable')
 
@@ -22,15 +23,32 @@ _logger = logging.getLogger('pixeltable')
 @pxt.udf(batch_size=4)
 def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0.5) -> Batch[dict]:
     """
-    Runs the specified YOLOX object detection model on an image.
+    Computes YOLOX object detections for the specified image. `model_id` should reference one of the models
+    defined in the [YOLOX documentation](https://github.com/Megvii-BaseDetection/YOLOX).
 
     YOLOX support is part of the `pixeltable.ext` package: long-term support is not guaranteed, and it is not
     intended for use in production applications.
 
-    Parameters:
-    - `model_id` - one of: `yolox_nano, `yolox_tiny`, `yolox_s`, `yolox_m`, `yolox_l`, `yolox_x`
-    - `threshold` - the threshold for object detection
+    __Requirements__:
+
+    - `pip install git+https://github.com/Megvii-BaseDetection/YOLOX`
+
+    Args:
+        model_id: one of: `yolox_nano`, `yolox_tiny`, `yolox_s`, `yolox_m`, `yolox_l`, `yolox_x`
+        threshold: the threshold for object detection
+
+    Returns:
+        A dictionary containing the output of the object detection model.
+
+    Examples:
+        Add a computed column that applies the model `yolox_m` to an existing
+        Pixeltable column `tbl.image` of the table `tbl`:
+
+        >>> tbl['detections'] = yolox(tbl.image, model_id='yolox_m', threshold=0.8)
     """
+    import torch
+    from yolox.utils import postprocess
+
     model, exp = _lookup_model(model_id, 'cpu')
     image_tensors = list(_images_to_tensors(images, exp))
     batch_tensor = torch.stack(image_tensors)
@@ -58,6 +76,21 @@ def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0
 
 @pxt.udf
 def yolo_to_coco(detections: dict) -> list:
+    """
+    Converts the output of a YOLOX object detection model to COCO format.
+
+    Args:
+        detections: The output of a YOLOX object detection model, as returned by `yolox`.
+
+    Returns:
+        A dictionary containing the data from `detections`, converted to COCO format.
+
+    Examples:
+        Add a computed column that converts the output `tbl.detections` to COCO format, where `tbl.image`
+        is the image for which detections were computed:
+
+        >>> tbl['detections_coco'] = yolo_to_coco(tbl.detections)
+    """
     bboxes, labels = detections['bboxes'], detections['labels']
     num_annotations = len(detections['bboxes'])
     assert num_annotations == len(detections['labels'])
@@ -72,14 +105,21 @@ def yolo_to_coco(detections: dict) -> list:
     return result
 
 
-def _images_to_tensors(images: Iterable[PIL.Image.Image], exp: Exp) -> Iterator[torch.Tensor]:
+def _images_to_tensors(images: Iterable[PIL.Image.Image], exp: 'Exp') -> Iterator['torch.Tensor']:
+    import torch
+    from yolox.data import ValTransform
+
+    _val_transform = ValTransform(legacy=False)
     for image in images:
         image = normalize_image_mode(image)
         image_transform, _ = _val_transform(np.array(image), None, exp.test_size)
         yield torch.from_numpy(image_transform)
 
 
-def _lookup_model(model_id: str, device: str) -> (YOLOX, Exp):
+def _lookup_model(model_id: str, device: str) -> tuple['YOLOX', 'Exp']:
+    import torch
+    from yolox.exp import get_exp
+
     key = (model_id, device)
     if key in _model_cache:
         return _model_cache[key]
@@ -105,5 +145,11 @@ def _lookup_model(model_id: str, device: str) -> (YOLOX, Exp):
     return model, exp
 
 
-_model_cache = {}
-_val_transform = ValTransform(legacy=False)
+_model_cache: dict[tuple[str, str], tuple['YOLOX', 'Exp']] = {}
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/func/aggregate_function.py

@@ -1,16 +1,18 @@
 from __future__ import annotations
 
 import abc
-import importlib
 import inspect
-from typing import Optional, Any, Type, List, Dict, Callable
-import itertools
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
+
 from .function import Function
-from .signature import Signature, Parameter
 from .globals import validate_symbol_path
+from .signature import Parameter, Signature
+
+if TYPE_CHECKING:
+    import pixeltable
 
 
 class Aggregator(abc.ABC):
@@ -40,6 +42,7 @@ class AggregateFunction(Function):
         self.requires_order_by = requires_order_by
         self.allows_std_agg = allows_std_agg
         self.allows_window = allows_window
+        self.__doc__ = aggregator_class.__doc__
 
         # our signature is the signature of 'update', but without self,
         # plus the parameters of 'init' as keyword-only parameters
@@ -135,6 +138,9 @@ class AggregateFunction(Function):
                     f'expression'
                 )
 
+    def __repr__(self) -> str:
+        return f'<Pixeltable Aggregator {self.name}>'
+
 
 def uda(
         *,

pixeltable/func/callable_function.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import inspect
-from typing import Optional, Callable, Tuple, Any
+from typing import Any, Callable, Optional
 from uuid import UUID
 
 import cloudpickle
@@ -19,13 +19,21 @@ class CallableFunction(Function):
     """
 
     def __init__(
-            self, signature: Signature, py_fn: Callable, self_path: Optional[str] = None,
-            self_name: Optional[str] = None, batch_size: Optional[int] = None):
+        self,
+        signature: Signature,
+        py_fn: Callable,
+        self_path: Optional[str] = None,
+        self_name: Optional[str] = None,
+        batch_size: Optional[int] = None,
+        is_method: bool = False,
+        is_property: bool = False
+    ):
         assert py_fn is not None
         self.py_fn = py_fn
         self.self_name = self_name
         self.batch_size = batch_size
-        super().__init__(signature, self_path=self_path)
+        self.__doc__ = py_fn.__doc__
+        super().__init__(signature, self_path=self_path, is_method=is_method, is_property=is_property)
 
     @property
     def is_batched(self) -> bool:
@@ -77,6 +85,7 @@ class CallableFunction(Function):
     def _as_dict(self) -> dict:
         if self.self_path is None:
             # this is not a module function
+            assert not self.is_method and not self.is_property
             from .function_registry import FunctionRegistry
             id = FunctionRegistry.get().create_stored_function(self)
             return {'id': id.hex}
@@ -113,3 +122,6 @@ class CallableFunction(Function):
                         f'{self.display_name}(): '
                         f'parameter {param.name} must be a constant value, not a Pixeltable expression'
                     )
+
+    def __repr__(self) -> str:
+        return f'<Pixeltable UDF {self.name}>'

pixeltable/func/expr_template_function.py

@@ -56,7 +56,7 @@ class ExprTemplateFunction(Function):
             arg_exprs[param_expr] = arg_expr
         result = result.substitute(arg_exprs)
         import pixeltable.exprs as exprs
-        assert not result.contains(exprs.Variable)
+        assert not result._contains(exprs.Variable)
         return result
 
     def exec(self, *args: Any, **kwargs: Any) -> Any:

pixeltable/func/function.py

@@ -3,10 +3,12 @@ from __future__ import annotations
 import abc
 import importlib
 import inspect
-from typing import Optional, Dict, Any, Tuple, Callable
+from typing import Any, Callable, Dict, Optional, Tuple
 
 import pixeltable
+import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
+
 from .globals import resolve_symbol
 from .signature import Signature
 
@@ -19,9 +21,13 @@ class Function(abc.ABC):
     via the member self_path.
     """
 
-    def __init__(self, signature: Signature, self_path: Optional[str] = None):
+    def __init__(self, signature: Signature, self_path: Optional[str] = None, is_method: bool = False, is_property: bool = False):
+        # Check that stored functions cannot be declared using `is_method` or `is_property`:
+        assert not ((is_method or is_property) and self_path is None)
         self.signature = signature
         self.self_path = self_path  # fully-qualified path to self
+        self.is_method = is_method
+        self.is_property = is_property
         self._conditional_return_type: Optional[Callable[..., ts.ColumnType]] = None
 
     @property
@@ -38,6 +44,10 @@ class Function(abc.ABC):
             return self.self_path[len(ptf_prefix):]
         return self.self_path
 
+    @property
+    def arity(self) -> int:
+        return len(self.signature.parameters)
+
     def help_str(self) -> str:
         return self.display_name + str(self.signature)
 

pixeltable/func/function_registry.py

@@ -4,11 +4,9 @@ import dataclasses
 import importlib
 import logging
 import sys
-import types
-from typing import Optional, Dict, List, Tuple
+from typing import Optional, Dict, List
 from uuid import UUID
 
-import cloudpickle
 import sqlalchemy as sql
 
 import pixeltable.env as env
@@ -36,6 +34,7 @@ class FunctionRegistry:
     def __init__(self):
         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]] = {}
 
     def clear_cache(self) -> None:
         """
@@ -66,7 +65,16 @@ class FunctionRegistry:
     #             self.module_fns[fn_path] = obj
 
     def register_function(self, fqn: str, fn: Function) -> None:
+        if fqn in self.module_fns:
+            raise excs.Error(f'A UDF with that name already exists: {fqn}')
         self.module_fns[fqn] = fn
+        if fn.is_method or fn.is_property:
+            base_type = fn.signature.parameters_by_pos[0].col_type.type_enum
+            if base_type not in self.type_methods:
+                self.type_methods[base_type] = {}
+            if fn.name in self.type_methods[base_type]:
+                raise excs.Error(f'Duplicate method name for type {base_type}: {fn.name}')
+            self.type_methods[base_type][fn.name] = fn
 
     def list_functions(self) -> List[Function]:
         # retrieve Function.Metadata data for all existing stored functions from store directly
@@ -127,12 +135,21 @@ class FunctionRegistry:
     #         assert fqn in self.module_fns, f'{fqn} not found'
     #         return self.module_fns[fqn]
 
-    def get_type_methods(self, name: str, base_type: ts.ColumnType.Type) -> List[Function]:
-        return [
-            fn for fn in self.module_fns.values()
-            if fn.self_path is not None and fn.self_path.endswith('.' + name) \
-               and fn.signature.parameters_by_pos[0].col_type.type_enum == base_type
-        ]
+    def get_type_methods(self, base_type: ts.ColumnType.Type) -> list[Function]:
+        """
+        Get a list of all methods (and properties) registered for a given base type.
+        """
+        if base_type in self.type_methods:
+            return list(self.type_methods[base_type].values())
+        return []
+
+    def lookup_type_method(self, base_type: ts.ColumnType.Type, name: str) -> Optional[Function]:
+        """
+        Look up a method (or property) by name for a given base type. If no such method is registered, return None.
+        """
+        if base_type in self.type_methods and name in self.type_methods[base_type]:
+            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: