pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/metadata/schema.py

@@ -1,35 +1,38 @@
 import dataclasses
+import types
 import typing
 import uuid
-from typing import Any, Optional, TypeVar, Union, get_type_hints
+from enum import Enum
+from typing import Any, TypeVar, Union, get_type_hints
 
 import sqlalchemy as sql
-import sqlalchemy.orm as orm
-from sqlalchemy import BigInteger, ForeignKey, Integer, LargeBinary
+from sqlalchemy import BigInteger, ForeignKey, Integer, LargeBinary, orm
 from sqlalchemy.dialects.postgresql import JSONB, UUID
-from sqlalchemy.orm import declarative_base
 from sqlalchemy.orm.decl_api import DeclarativeMeta
 
+from ..catalog.update_status import UpdateStatus
+
 # Base has to be marked explicitly as a type, in order to be used elsewhere as a type hint. But in addition to being
 # a type, it's also a `DeclarativeMeta`. The following pattern enables us to expose both `Base` and `Base.metadata`
 # outside of the module in a typesafe way.
-Base: type = declarative_base()
+Base: type = orm.declarative_base()
 assert isinstance(Base, DeclarativeMeta)
 base_metadata = Base.metadata
 
 T = TypeVar('T')
 
-def md_from_dict(data_class_type: type[T], data: Any) -> T:
+
+def md_from_dict(type_: type[T], data: Any) -> T:
     """Re-instantiate a dataclass instance that contains nested dataclasses from a dict."""
-    if dataclasses.is_dataclass(data_class_type):
-        fieldtypes = {f: t for f, t in get_type_hints(data_class_type).items()}
-        return data_class_type(**{f: md_from_dict(fieldtypes[f], data[f]) for f in data})  # type: ignore[return-value]
+    if dataclasses.is_dataclass(type_):
+        fieldtypes = get_type_hints(type_)
+        return type_(**{f: md_from_dict(fieldtypes[f], data[f]) for f in data})
 
-    origin = typing.get_origin(data_class_type)
+    origin = typing.get_origin(type_)
     if origin is not None:
-        type_args = typing.get_args(data_class_type)
-        if origin is Union and type(None) in type_args:
-            # Handling Optional types
+        type_args = typing.get_args(type_)
+        if (origin is Union or origin is types.UnionType) and type(None) in type_args:
+            # handling T | None, T | None
             non_none_args = [arg for arg in type_args if arg is not type(None)]
             assert len(non_none_args) == 1
             return md_from_dict(non_none_args[0], data) if data is not None else None
@@ -42,11 +45,19 @@ def md_from_dict(data_class_type: type[T], data: Any) -> T:
         elif origin is tuple:
             return tuple(md_from_dict(arg_type, elem) for arg_type, elem in zip(type_args, data))  # type: ignore[return-value]
         else:
-            assert False
+            raise AssertionError(origin)
+    elif isinstance(type_, type) and issubclass(type_, Enum):
+        return type_(data)
     else:
         return data
 
 
+def _md_dict_factory(data: list[tuple[str, Any]]) -> dict:
+    """Use this to serialize <>Md instances with dataclasses.asdict()"""
+    # serialize enums to their values
+    return {k: v.value if isinstance(v, Enum) else v for k, v in data}
+
+
 # structure of the stored metadata:
 # - each schema entity that grows somehow proportionally to the data (# of output_rows, total insert operations,
 #   number of schema changes) gets its own table
@@ -56,6 +67,7 @@ def md_from_dict(data_class_type: type[T], data: Any) -> T:
 #   schema easier (the goal is not to have to rely on some schema migration framework; if that breaks for some user,
 #   it would be very difficult to patch up)
 
+
 @dataclasses.dataclass
 class SystemInfoMd:
     schema_version: int
@@ -63,7 +75,9 @@ class SystemInfoMd:
 
 class SystemInfo(Base):
     """A single-row table that contains system-wide metadata."""
+
     __tablename__ = 'systeminfo'
+
     dummy = sql.Column(Integer, primary_key=True, default=0, nullable=False)
     md = sql.Column(JSONB, nullable=False)  # SystemInfoMd
 
@@ -71,14 +85,22 @@ class SystemInfo(Base):
 @dataclasses.dataclass
 class DirMd:
     name: str
+    user: str | None
+    additional_md: dict[str, Any]  # deprecated
 
 
 class Dir(Base):
     __tablename__ = 'dirs'
 
-    id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, nullable=False)
+    id: orm.Mapped[uuid.UUID] = orm.mapped_column(
+        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, nullable=False
+    )
     parent_id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), ForeignKey('dirs.id'), nullable=True)
-    md = sql.Column(JSONB, nullable=False)
+    md: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False)  # DirMd
+    additional_md: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False, default=dict)
+
+    # used to force acquisition of an X-lock via an Update stmt
+    lock_dummy: orm.Mapped[int] = orm.mapped_column(BigInteger, nullable=True)
 
 
 @dataclasses.dataclass
@@ -89,20 +111,24 @@ class ColumnMd:
     - when a column was added/dropped, which is needed to GC unreachable storage columns
       (a column that was added after table snapshot n and dropped before table snapshot n+1 can be removed
       from the stored table).
-     """
+    """
+
     id: int
     schema_version_add: int
-    schema_version_drop: Optional[int]
+    schema_version_drop: int | None
     col_type: dict
 
     # if True, is part of the primary key
     is_pk: bool
 
     # if set, this is a computed column
-    value_expr: Optional[dict]
+    value_expr: dict | None
 
     # if True, the column is present in the stored table
-    stored: Optional[bool]
+    stored: bool | None
+
+    # If present, the URI for the destination for column values
+    destination: str | None = None
 
 
 @dataclasses.dataclass
@@ -110,37 +136,71 @@ class IndexMd:
     """
     Metadata needed to instantiate an EmbeddingIndex
     """
+
     id: int
     name: str
+    indexed_col_tbl_id: str  # UUID of the table (as string) that contains column being indexed
     indexed_col_id: int  # column being indexed
     index_val_col_id: int  # column holding the values to be indexed
     index_val_undo_col_id: int  # column holding index values for deleted rows
     schema_version_add: int
-    schema_version_drop: Optional[int]
+    schema_version_drop: int | None
     class_fqn: str
     init_args: dict[str, Any]
 
 
+# a stored table version path is a list of (table id as str, effective table version)
+TableVersionPath = list[tuple[str, int | None]]
+
+
 @dataclasses.dataclass
 class ViewMd:
     is_snapshot: bool
+    include_base_columns: bool
 
     # (table id, version); for mutable views, all versions are None
-    base_versions: list[tuple[str, Optional[int]]]
+    base_versions: TableVersionPath
 
     # filter predicate applied to the base table; view-only
-    predicate: Optional[dict[str, Any]]
+    predicate: dict[str, Any] | None
+
+    # sampling predicate applied to the base table; view-only
+    sample_clause: dict[str, Any] | None
 
     # ComponentIterator subclass; only for component views
-    iterator_class_fqn: Optional[str]
+    iterator_class_fqn: str | None
 
     # args to pass to the iterator class constructor; only for component views
-    iterator_args: Optional[dict[str, Any]]
+    iterator_args: dict[str, Any] | None
+
+
+class TableState(Enum):
+    """The operational state of the table"""
+
+    LIVE = 0
+    ROLLFORWARD = 1  # finalizing pending table ops
+    ROLLBACK = 2  # rolling back pending table ops
+
+
+class TableStatement(Enum):
+    """The top-level DDL/DML operation (corresponding to a statement in SQL; not: a TableOp) currently being executed"""
+
+    CREATE_TABLE = 0
+    CREATE_VIEW = 1
+    DROP_TABLE = 2
+    ADD_COLUMNS = 3
+    DROP_COLUMNS = 4
+    ADD_INDEX = 5
+    DROP_INDEX = 6
 
 
 @dataclasses.dataclass
 class TableMd:
+    tbl_id: str  # uuid.UUID
     name: str
+    is_replica: bool
+
+    user: str | None
 
     # monotonically increasing w/in Table for both data and schema changes, starting at 0
     current_version: int
@@ -154,13 +214,52 @@ class TableMd:
     # - every row is assigned a unique and immutable rowid on insertion
     next_row_id: int
 
+    # sequence number to track changes in the set of mutable views of this table (ie, this table = the view base)
+    # - incremented for each add/drop of a mutable view
+    # - only maintained for mutable tables
+    # TODO: replace with mutable_views: list[UUID] to help with debugging
+    view_sn: int
+
     # Metadata format for external stores:
     # {'class': 'pixeltable.io.label_studio.LabelStudioProject', 'md': {'project_id': 3}}
     external_stores: list[dict[str, Any]]
 
     column_md: dict[int, ColumnMd]  # col_id -> ColumnMd
     index_md: dict[int, IndexMd]  # index_id -> IndexMd
-    view_md: Optional[ViewMd]
+    view_md: ViewMd | None
+    # TODO: Remove additional_md from this and other Md dataclasses (and switch to using the separate additional_md
+    #     column in all cases)
+    additional_md: dict[str, Any]  # deprecated
+
+    # deprecated
+    has_pending_ops: bool = False
+
+    tbl_state: TableState = TableState.LIVE
+    pending_stmt: TableStatement | None = None
+
+    @property
+    def is_snapshot(self) -> bool:
+        return self.view_md is not None and self.view_md.is_snapshot
+
+    @property
+    def is_mutable(self) -> bool:
+        return not self.is_snapshot and not self.is_replica
+
+    @property
+    def is_pure_snapshot(self) -> bool:
+        return (
+            self.view_md is not None
+            and self.view_md.is_snapshot
+            and self.view_md.sample_clause is None
+            and self.view_md.predicate is None
+            and len(self.column_md) == 0
+        )
+
+    @property
+    def ancestors(self) -> TableVersionPath:
+        if self.view_md is None:
+            return []
+        return self.view_md.base_versions
 
 
 class Table(Base):
@@ -170,28 +269,46 @@ class Table(Base):
     Views are in essence a subclass of tables, because they also store materialized columns. The differences are:
     - views have a base, which is either a (live) table or a snapshot
     - views can have a filter predicate
+
+    dir_id: NULL for dropped tables
     """
+
     __tablename__ = 'tables'
 
     MAX_VERSION = 9223372036854775807  # 2^63 - 1
 
     id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), primary_key=True, nullable=False)
-    dir_id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), ForeignKey('dirs.id'), nullable=False)
-    md = sql.Column(JSONB, nullable=False)  # TableMd
+    dir_id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), ForeignKey('dirs.id'), nullable=True)
+    md: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False)  # TableMd
+    additional_md: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False, default=dict)
+
+    # used to force acquisition of an X-lock via an Update stmt
+    lock_dummy: orm.Mapped[int] = orm.mapped_column(BigInteger, nullable=True)
 
 
 @dataclasses.dataclass
-class TableVersionMd:
+class VersionMd:
+    tbl_id: str  # uuid.UUID
     created_at: float  # time.time()
     version: int
     schema_version: int
+    user: str | None = None  # User that created this version
+    update_status: UpdateStatus | None = None  # UpdateStatus of the change that created this version
+    # A version fragment cannot be queried or instantiated via get_table(). A fragment represents a version of a
+    # replica table that has incomplete data, and exists only to provide base table support for a dependent view.
+    is_fragment: bool = False
+    additional_md: dict[str, Any] = dataclasses.field(default_factory=dict)  # deprecated
 
 
 class TableVersion(Base):
     __tablename__ = 'tableversions'
-    tbl_id = sql.Column(UUID(as_uuid=True), ForeignKey('tables.id'), primary_key=True, nullable=False)
-    version = sql.Column(BigInteger, primary_key=True, nullable=False)
-    md = sql.Column(JSONB, nullable=False)  # TableVersionMd
+
+    tbl_id: orm.Mapped[uuid.UUID] = orm.mapped_column(
+        UUID(as_uuid=True), ForeignKey('tables.id'), primary_key=True, nullable=False
+    )
+    version: orm.Mapped[int] = orm.mapped_column(BigInteger, primary_key=True, nullable=False)
+    md: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False)
+    additional_md: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False, default=dict)
 
 
 @dataclasses.dataclass
@@ -199,21 +316,24 @@ class SchemaColumn:
     """
     Records the versioned metadata of a column.
     """
+
     pos: int
     name: str
 
     # media validation strategy of this particular media column; if not set, TableMd.media_validation applies
     # stores column.MediaValiation.name.lower()
-    media_validation: Optional[str]
+    media_validation: str | None
 
 
 @dataclasses.dataclass
-class TableSchemaVersionMd:
+class SchemaVersionMd:
     """
     Records all versioned table metadata.
     """
+
+    tbl_id: str  # uuid.UUID
     schema_version: int
-    preceding_schema_version: Optional[int]
+    preceding_schema_version: int | None
     columns: dict[int, SchemaColumn]  # col_id -> SchemaColumn
     num_retained_versions: int
     comment: str
@@ -221,15 +341,35 @@ class TableSchemaVersionMd:
     # default validation strategy for any media column of this table
     # stores column.MediaValiation.name.lower()
     media_validation: str
+    additional_md: dict[str, Any]  # deprecated
 
 
 # versioning: each table schema change results in a new record
 class TableSchemaVersion(Base):
     __tablename__ = 'tableschemaversions'
 
-    tbl_id = sql.Column(UUID(as_uuid=True), ForeignKey('tables.id'), primary_key=True, nullable=False)
-    schema_version = sql.Column(BigInteger, primary_key=True, nullable=False)
-    md = sql.Column(JSONB, nullable=False)  # TableSchemaVersionMd
+    tbl_id: orm.Mapped[uuid.UUID] = orm.mapped_column(
+        UUID(as_uuid=True), ForeignKey('tables.id'), primary_key=True, nullable=False
+    )
+    schema_version: orm.Mapped[int] = orm.mapped_column(BigInteger, primary_key=True, nullable=False)
+    md: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False)  # TableSchemaVersionMd
+    additional_md: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False, default=dict)
+
+
+class PendingTableOp(Base):
+    """
+    Table operation that needs to be completed before the table can be used.
+
+    Operations need to be completed in order of increasing seq_num.
+    """
+
+    __tablename__ = 'pendingtableops'
+
+    tbl_id: orm.Mapped[uuid.UUID] = orm.mapped_column(
+        UUID(as_uuid=True), ForeignKey('tables.id'), primary_key=True, nullable=False
+    )
+    op_sn: orm.Mapped[int] = orm.mapped_column(Integer, primary_key=True, nullable=False)  # catalog.TableOp.op_sn
+    op: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False)  # catalog.TableOp
 
 
 @dataclasses.dataclass
@@ -249,9 +389,12 @@ class Function(Base):
     We store the Python version under which a Function was created (and the callable pickled) in order to warn
     against version mismatches.
     """
+
     __tablename__ = 'functions'
 
-    id = sql.Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, nullable=False)
-    dir_id = sql.Column(UUID(as_uuid=True), ForeignKey('dirs.id'), nullable=True)
-    md = sql.Column(JSONB, nullable=False)  # FunctionMd
-    binary_obj = sql.Column(LargeBinary, nullable=True)
+    id: orm.Mapped[uuid.UUID] = orm.mapped_column(
+        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, nullable=False
+    )
+    dir_id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), ForeignKey('dirs.id'), nullable=True)
+    md: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False)  # FunctionMd
+    binary_obj: orm.Mapped[bytes | None] = orm.mapped_column(LargeBinary, nullable=True)

pixeltable/metadata/utils.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from pixeltable.metadata import schema
+
+
+class MetadataUtils:
+    @classmethod
+    def _diff_md(
+        cls, old_md: dict[int, schema.SchemaColumn] | None, new_md: dict[int, schema.SchemaColumn] | None
+    ) -> str:
+        """Return a string reporting the differences in a specific entry in two dictionaries
+
+        Results are formatted as follows:
+        - If `old_md` is `None`, returns 'Initial Version'.
+        - If `old_md` and `new_md` are the same, returns an empty string.
+        - If there are additions, changes, or deletions, returns a string summarizing the changes.
+        """
+        assert new_md is not None
+        if old_md is None:
+            return 'Initial Version'
+        if old_md == new_md:
+            return ''
+        added = {k: v.name for k, v in new_md.items() if k not in old_md}
+        changed = {
+            k: f'{old_md[k].name!r} to {v.name!r}'
+            for k, v in new_md.items()
+            if k in old_md and old_md[k].name != v.name
+        }
+        deleted = {k: v.name for k, v in old_md.items() if k not in new_md}
+        if len(added) == 0 and len(changed) == 0 and len(deleted) == 0:
+            return ''
+        # Format the result
+        t = []
+        if len(added) > 0:
+            t.append('Added: ' + ', '.join(added.values()))
+        if len(changed) > 0:
+            t.append('Renamed: ' + ', '.join(changed.values()))
+        if len(deleted) > 0:
+            t.append('Deleted: ' + ', '.join(deleted.values()))
+        r = ', '.join(t)
+        return r
+
+    @classmethod
+    def _create_md_change_dict(cls, md_list: list[tuple[int, dict[int, schema.SchemaColumn]]] | None) -> dict[int, str]:
+        """Return a dictionary of schema changes by version
+        Args:
+            md_list: a list of tuples, each containing a version number and a metadata dictionary.
+        """
+        r: dict[int, str] = {}
+        if md_list is None or len(md_list) == 0:
+            return r
+
+        # Sort the list in place by version number
+        md_list.sort()
+
+        first_retrieved_version = md_list[0][0]
+        if first_retrieved_version == 0:
+            prev_md = None
+            prev_ver = -1
+            start = 0
+        else:
+            prev_md = md_list[0][1]
+            prev_ver = first_retrieved_version
+            start = 1
+
+        for ver, curr_md in md_list[start:]:
+            if ver == prev_ver:
+                continue
+            assert ver > prev_ver
+            tf = cls._diff_md(prev_md, curr_md)
+            if tf != '':
+                r[ver] = tf
+            prev_md = curr_md
+        return r

pixeltable/mypy/__init__.py

@@ -0,0 +1,3 @@
+from .mypy_plugin import plugin
+
+__all__ = ['plugin']

pixeltable/mypy/mypy_plugin.py

@@ -0,0 +1,123 @@
+from typing import Callable, ClassVar
+
+from mypy import nodes
+from mypy.plugin import AnalyzeTypeContext, ClassDefContext, FunctionContext, MethodSigContext, Plugin
+from mypy.plugins.common import add_attribute_to_class, add_method_to_class
+from mypy.types import AnyType, FunctionLike, Instance, NoneType, Type, TypeOfAny
+
+import pixeltable as pxt
+from pixeltable import exprs
+
+
+class PxtPlugin(Plugin):
+    __UDA_FULLNAME = f'{pxt.uda.__module__}.{pxt.uda.__name__}'
+    __ARRAY_GETITEM_FULLNAME = f'{pxt.Array.__module__}.{pxt.Array.__name__}.__class_getitem__'
+    __ADD_COLUMN_FULLNAME = f'{pxt.Table.__module__}.{pxt.Table.__name__}.{pxt.Table.add_column.__name__}'
+    __ADD_COMPUTED_COLUMN_FULLNAME = (
+        f'{pxt.Table.__module__}.{pxt.Table.__name__}.{pxt.Table.add_computed_column.__name__}'
+    )
+    __TYPE_MAP: ClassVar[dict] = {
+        pxt.Json: 'typing.Any',
+        pxt.Array: 'numpy.ndarray',
+        pxt.Image: 'PIL.Image.Image',
+        pxt.Video: 'builtins.str',
+        pxt.Audio: 'builtins.str',
+        pxt.Document: 'builtins.str',
+    }
+    __FULLNAME_MAP: ClassVar[dict] = {f'{k.__module__}.{k.__name__}': v for k, v in __TYPE_MAP.items()}
+
+    def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] | None:
+        return adjust_uda_type
+
+    def get_type_analyze_hook(self, fullname: str) -> Callable[[AnalyzeTypeContext], Type] | None:
+        if fullname in self.__FULLNAME_MAP:
+            subst_name = self.__FULLNAME_MAP[fullname]
+            return lambda ctx: adjust_pxt_type(ctx, subst_name)
+        return None
+
+    def get_method_signature_hook(self, fullname: str) -> Callable[[MethodSigContext], FunctionLike] | None:
+        if fullname in (self.__ADD_COLUMN_FULLNAME, self.__ADD_COMPUTED_COLUMN_FULLNAME):
+            return adjust_kwargs
+        return None
+
+    def get_class_decorator_hook_2(self, fullname: str) -> Callable[[ClassDefContext], bool] | None:
+        if fullname == self.__UDA_FULLNAME:
+            return adjust_uda_methods
+        return None
+
+
+def plugin(version: str) -> type:
+    return PxtPlugin
+
+
+_AGGREGATOR_FULLNAME = f'{pxt.Aggregator.__module__}.{pxt.Aggregator.__name__}'
+_FN_CALL_FULLNAME = f'{exprs.Expr.__module__}.{exprs.Expr.__name__}'
+
+
+def adjust_uda_type(ctx: FunctionContext) -> Type:
+    """
+    Mypy doesn't understand that a class with a @uda decorator isn't actually a class, so it assumes
+    that sum(expr), for example, actually returns an instance of sum. We correct this by changing the
+    return type of any subclass of `Aggregator` to `FunctionCall`.
+    """
+    ret_type = ctx.default_return_type
+    if isinstance(ret_type, Instance) and (
+        ret_type.type.fullname == _AGGREGATOR_FULLNAME
+        or any(base.type.fullname == _AGGREGATOR_FULLNAME for base in ret_type.type.bases)
+    ):
+        ret_type = AnyType(TypeOfAny.special_form)
+    return ret_type
+
+
+def adjust_pxt_type(ctx: AnalyzeTypeContext, subst_name: str) -> Type:
+    """
+    Replaces the special Pixeltable classes (such as pxt.Array) with their standard equivalents (such as np.ndarray).
+    """
+    if subst_name == 'typing.Any':
+        return AnyType(TypeOfAny.special_form)
+    return ctx.api.named_type(subst_name, [])
+
+
+def adjust_kwargs(ctx: MethodSigContext) -> FunctionLike:
+    """
+    Mypy has a "feature" where it will spit out multiple warnings if a method with signature
+    ```
+    def my_func(*, arg1: int, arg2: str, **kwargs: Expr)
+    ```
+    (for example) is called with bare kwargs:
+    ```
+    my_func(my_kwarg=value)
+    ```
+    This is a disaster for type-checking of add_column and add_computed_column. Here we adjust the signature so
+    that mypy thinks it is simply
+    ```
+    def my_func(**kwargs: Any)
+    ```
+    thereby avoiding any type-checking errors. For details, see: <https://github.com/python/mypy/issues/18481>
+    """
+    sig = ctx.default_signature
+    new_arg_names = sig.arg_names[-1:]
+    new_arg_types = [AnyType(TypeOfAny.special_form)]
+    new_arg_kinds = sig.arg_kinds[-1:]
+    return sig.copy_modified(arg_names=new_arg_names, arg_types=new_arg_types, arg_kinds=new_arg_kinds)
+
+
+def adjust_uda_methods(ctx: ClassDefContext) -> bool:
+    """
+    Mypy does not handle the `@pxt.uda` aggregator well; it continues to treat the decorated class as a class,
+    even though it has been replaced by an `AggregateFunction`. Here we add static methods to the class that
+    imitate various (instance) methods of `AggregateFunction` so that they can be properly type-checked.
+    """
+    list_type = ctx.api.named_type('builtins.list', [AnyType(TypeOfAny.special_form)])
+    fn_arg = nodes.Argument(nodes.Var('fn'), AnyType(TypeOfAny.special_form), None, nodes.ARG_POS)
+    args_arg = nodes.Argument(nodes.Var('args'), AnyType(TypeOfAny.special_form), None, nodes.ARG_STAR)
+    kwargs_arg = nodes.Argument(nodes.Var('kwargs'), AnyType(TypeOfAny.special_form), None, nodes.ARG_STAR2)
+    add_method_to_class(ctx.api, ctx.cls, '__init__', args=[args_arg, kwargs_arg], return_type=NoneType())
+    add_method_to_class(
+        ctx.api, ctx.cls, 'to_sql', args=[fn_arg], return_type=AnyType(TypeOfAny.special_form), is_staticmethod=True
+    )
+    add_method_to_class(
+        ctx.api, ctx.cls, 'overload', args=[fn_arg], return_type=AnyType(TypeOfAny.special_form), is_staticmethod=True
+    )
+    add_attribute_to_class(ctx.api, ctx.cls, 'signatures', typ=list_type, is_classvar=True)
+    return True