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

pixeltable/functions/globals.py

@@ -23,7 +23,7 @@ T = typing.TypeVar('T')
 class sum(func.Aggregator, typing.Generic[T]):
     """Sums the selected integers or floats."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.sum: T = None
 
     def update(self, val: T) -> None:
@@ -67,7 +67,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     ),
 )
 class count(func.Aggregator, typing.Generic[T]):
-    def __init__(self):
+    def __init__(self) -> None:
         self.count = 0
 
     def update(self, val: T) -> None:
@@ -88,7 +88,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     type_substitutions=tuple({T: Optional[t]} for t in (str, int, float, bool, ts.Timestamp)),  # type: ignore[misc]
 )
 class min(func.Aggregator, typing.Generic[T]):
-    def __init__(self):
+    def __init__(self) -> None:
         self.val: T = None
 
     def update(self, val: T) -> None:
@@ -118,7 +118,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     type_substitutions=tuple({T: Optional[t]} for t in (str, int, float, bool, ts.Timestamp)),  # type: ignore[misc]
 )
 class max(func.Aggregator, typing.Generic[T]):
-    def __init__(self):
+    def __init__(self) -> None:
         self.val: T = None
 
     def update(self, val: T) -> None:
@@ -143,7 +143,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
 
 @func.uda(type_substitutions=({T: Optional[int]}, {T: Optional[float]}))  # type: ignore[misc]
 class mean(func.Aggregator, typing.Generic[T]):
-    def __init__(self):
+    def __init__(self) -> None:
         self.sum: T = None
         self.count = 0
 
@@ -182,5 +182,5 @@ def map(expr: exprs.Expr, fn: Callable[[exprs.Expr], Any]) -> exprs.Expr:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/huggingface.py

@@ -493,5 +493,5 @@ _processor_cache: dict[tuple[str, Callable], Any] = {}
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/image.py

@@ -426,5 +426,5 @@ def mode(self: PIL.Image.Image) -> str:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/json.py

@@ -37,5 +37,5 @@ class make_list(pxt.Aggregator):
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/llama_cpp.py

@@ -100,5 +100,5 @@ _IS_GPU_AVAILABLE: Optional[bool] = None
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/math.py

@@ -63,5 +63,5 @@ def _(self: sql.ColumnElement, digits: Optional[sql.ColumnElement] = None) -> sq
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/mistralai.py

@@ -193,5 +193,5 @@ def _opt(arg: Optional[_T]) -> Union[_T, 'mistralai.types.basemodel.Unset']:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/ollama.py

@@ -129,5 +129,5 @@ def embed(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/openai.py

@@ -122,7 +122,7 @@ class OpenAIRateLimitsInfo(env.RateLimitsInfo):
 _header_duration_pattern = re.compile(r'(?:(\d+)d)?(?:(\d+)h)?(?:(\d+)ms)|(?:(\d+)m)?(?:([\d.]+)s)?')
 
 
-def _parse_header_duration(duration_str):
+def _parse_header_duration(duration_str: str) -> datetime.timedelta:
     match = _header_duration_pattern.match(duration_str)
     if not match:
         raise ValueError('Invalid duration format')
@@ -837,5 +837,5 @@ def _opt(arg: _T) -> Union[_T, 'openai.NotGiven']:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/replicate.py

@@ -69,5 +69,5 @@ async def run(input: dict[str, Any], *, ref: str) -> dict[str, Any]:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/string.py

@@ -675,5 +675,5 @@ def zfill(self: str, width: int) -> str:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/timestamp.py

@@ -312,5 +312,5 @@ def posix_timestamp(self: datetime) -> float:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/together.py

@@ -298,5 +298,5 @@ async def image_generations(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/util.py

@@ -21,7 +21,7 @@ def normalize_image_mode(image: PIL.Image.Image) -> PIL.Image.Image:
     Converts grayscale images to 3-channel for compatibility with models that only work with
     multichannel input.
     """
-    if image.mode in {'1', 'L'}:
+    if image.mode in ('1', 'L'):
         return image.convert('RGB')
     if image.mode == 'LA':
         return image.convert('RGBA')

pixeltable/functions/video.py

@@ -143,7 +143,7 @@ def _get_metadata(path: str) -> dict:
 
 
 def __get_stream_metadata(stream: av.stream.Stream) -> dict:
-    if stream.type not in {'audio', 'video'}:
+    if stream.type not in ('audio', 'video'):
         return {'type': stream.type}  # Currently unsupported
 
     codec_context = stream.codec_context
@@ -190,5 +190,5 @@ def __get_stream_metadata(stream: av.stream.Stream) -> dict:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/vision.py

@@ -238,7 +238,7 @@ class mean_ap(pxt.Aggregator):
     - A `dict[int, float]` mapping each label class to an average precision (AP) value for that class.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.class_tpfp: dict[int, list[dict]] = defaultdict(list)
 
     def update(self, eval_dicts: list[dict]) -> None:
@@ -394,5 +394,5 @@ def draw_bounding_boxes(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/index/embedding_index.py

@@ -208,7 +208,18 @@ class EmbeddingIndex(IndexBase):
                 and len(sig.required_parameters) <= 1
                 and sig.parameters_by_pos[0].col_type.type_enum == expected_type
             ):
-                return resolved_fn
+                # We found a valid signature. Now, if it has more than one parameter, we need to transform it into a
+                # 1-ary function by fixing all the other parameters to their defaults. This is to ensure that
+                # conditional_return_type resolves correctly.
+                if len(sig.parameters) == 1:
+                    unary_fn = resolved_fn
+                else:
+                    assert all(sig.parameters_by_pos[i].has_default for i in range(1, len(sig.parameters)))
+                    defaults = {param.name: param.default for param in sig.parameters_by_pos[1:]}
+                    unary_fn = resolved_fn.using(**defaults)
+                assert not unary_fn.is_polymorphic
+                assert len(unary_fn.signature.parameters) == 1
+                return unary_fn
         return None
 
     @classmethod

pixeltable/io/__init__.py

@@ -1,3 +1,5 @@
+# ruff: noqa: F401
+
 from .datarows import import_json, import_rows
 from .external_store import ExternalStore, SyncStatus
 from .globals import create_label_studio_project, export_images_as_fo_dataset
@@ -5,10 +7,10 @@ from .hf_datasets import import_huggingface_dataset
 from .pandas import import_csv, import_excel, import_pandas
 from .parquet import export_parquet, import_parquet
 
-__default_dir = set(symbol for symbol in dir() if not symbol.startswith('_'))
+__default_dir = {symbol for symbol in dir() if not symbol.startswith('_')}
 __removed_symbols = {'globals', 'hf_datasets', 'pandas', 'parquet', 'datarows'}
-__all__ = sorted(list(__default_dir - __removed_symbols))
+__all__ = sorted(__default_dir - __removed_symbols)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/io/fiftyone.py

@@ -1,5 +1,5 @@
 import os
-from typing import Iterator, Optional, Union
+from typing import Any, Iterator, Optional, Union
 
 import fiftyone as fo  # type: ignore[import-untyped]
 import fiftyone.utils.data as foud  # type: ignore[import-untyped]
@@ -59,10 +59,9 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
         for exprs_, label_cls, default_name in label_categories:
             if exprs_ is None or isinstance(exprs_, dict):
                 continue
-            if isinstance(exprs_, exprs.Expr):
-                exprs_ = [exprs_]
-            assert isinstance(exprs_, list)
-            for expr in exprs_:
+            exprs_list = [exprs_] if isinstance(exprs_, exprs.Expr) else exprs_
+            assert isinstance(exprs_list, list)
+            for expr in exprs_list:
                 if default_name not in self.__labels:
                     name = default_name
                 else:
@@ -125,7 +124,7 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
             elif label_cls is fo.Detections:
                 label = fo.Detections(detections=self.__as_fo_detections(label_data))
             else:
-                assert False
+                raise AssertionError()
             labels[label_name] = label
 
         return file, metadata, labels
@@ -166,5 +165,5 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
     def get_dataset_info(self) -> dict:
         pass
 
-    def close(self, *args) -> None:
+    def close(self, *args: Any) -> None:
         pass

pixeltable/io/label_studio.py

@@ -5,16 +5,14 @@ import os
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Iterator, Literal, Optional, cast
-from xml.etree import ElementTree
+from xml.etree import ElementTree as ET
 
 import label_studio_sdk  # type: ignore[import-untyped]
 import PIL.Image
 from requests.exceptions import HTTPError
 
 import pixeltable as pxt
-import pixeltable.env as env
-import pixeltable.exceptions as excs
-from pixeltable import Column, Table
+from pixeltable import Column, Table, env, exceptions as excs
 from pixeltable.config import Config
 from pixeltable.exprs import ColumnRef, DataRow, Expr
 from pixeltable.io.external_store import Project, SyncStatus
@@ -140,7 +138,8 @@ class LabelStudioProject(Project):
             page += 1
         if unknown_task_count > 0:
             _logger.warning(
-                f'Skipped {unknown_task_count} unrecognized task(s) when syncing Label Studio project "{self.project_title}".'
+                f'Skipped {unknown_task_count} unrecognized task(s) when syncing '
+                f'Label Studio project {self.project_title!r}.'
             )
 
     def __update_tasks(self, t: Table, existing_tasks: dict[tuple, dict]) -> SyncStatus:
@@ -174,11 +173,11 @@ class LabelStudioProject(Project):
             # Send media to Label Studio by HTTP post.
             assert len(t_data_cols) == 1  # This was verified when the project was set up
             return self.__update_tasks_by_post(t, existing_tasks, t_data_cols[0], t_rl_cols, rl_info)
-        elif self.media_import_method == 'file' or self.media_import_method == 'url':
+        elif self.media_import_method in ('file', 'url'):
             # Send media to Label Studio by file reference (local file or URL).
             return self.__update_tasks_by_files(t, existing_tasks, t_data_cols, t_rl_cols, rl_info)
         else:
-            assert False
+            raise AssertionError()
 
     def __update_tasks_by_post(
         self,
@@ -227,7 +226,7 @@ class LabelStudioProject(Project):
                     )
                     for i in range(len(coco_annotations))
                 ]
-                _logger.debug(f'`predictions`: %s', predictions)
+                _logger.debug('`predictions`: {%s}', predictions)
                 self.project.create_predictions(predictions)
                 tasks_created += 1
 
@@ -358,7 +357,7 @@ class LabelStudioProject(Project):
     def __localpath_to_lspath(cls, localpath: str) -> str:
         # Transform the local path into Label Studio's bespoke path format.
         relpath = Path(localpath).relative_to(Config.get().home)
-        return f'/data/local-files/?d={str(relpath)}'
+        return f'/data/local-files/?d={relpath}'
 
     def __delete_stale_tasks(
         self, existing_tasks: dict[tuple, dict], row_ids_in_pxt: set[tuple], tasks_created: int
@@ -405,7 +404,8 @@ class LabelStudioProject(Project):
         updates = [{'_rowid': rowid, local_annotations_col.name: ann} for rowid, ann in annotations.items()]
         if len(updates) > 0:
             _logger.info(
-                f'Updating table `{t._name}`, column `{local_annotations_col.name}` with {len(updates)} total annotations.'
+                f'Updating table {t._name!r}, column {local_annotations_col.name!r} '
+                f'with {len(updates)} total annotations.'
             )
             # batch_update currently doesn't propagate from views to base tables. As a workaround, we call
             # batch_update on the actual ancestor table that holds the annotations column.
@@ -451,7 +451,7 @@ class LabelStudioProject(Project):
         Parses a Label Studio XML config, extracting the names and Pixeltable types of
         all input variables.
         """
-        root: ElementTree.Element = ElementTree.fromstring(xml_config)
+        root: ET.Element = ET.fromstring(xml_config)
         if root.tag.lower() != 'view':
             raise excs.Error('Root of Label Studio config must be a `View`')
         config = _LabelStudioConfig(
@@ -461,7 +461,7 @@ class LabelStudioProject(Project):
         return config
 
     @classmethod
-    def __parse_data_keys_config(cls, root: ElementTree.Element) -> dict[str, '_DataKey']:
+    def __parse_data_keys_config(cls, root: ET.Element) -> dict[str, '_DataKey']:
         """Parses the data keys from a Label Studio XML config."""
         config: dict[str, '_DataKey'] = {}
         for element in root:
@@ -477,7 +477,7 @@ class LabelStudioProject(Project):
         return config
 
     @classmethod
-    def __parse_rectangle_labels_config(cls, root: ElementTree.Element) -> dict[str, '_RectangleLabel']:
+    def __parse_rectangle_labels_config(cls, root: ET.Element) -> dict[str, '_RectangleLabel']:
         """Parses the RectangleLabels from a Label Studio XML config."""
         config: dict[str, '_RectangleLabel'] = {}
         for element in root:
@@ -534,7 +534,7 @@ class LabelStudioProject(Project):
         _label_studio_client().delete_project(self.project_id)
         env.Env.get().console_logger.info(f'Deleted Label Studio project: {title}')
 
-    def __eq__(self, other) -> bool:
+    def __eq__(self, other: object) -> bool:
         return isinstance(other, LabelStudioProject) and self.project_id == other.project_id
 
     def __hash__(self) -> int:
@@ -576,7 +576,7 @@ class LabelStudioProject(Project):
                 local_annotations_column = ANNOTATIONS_COLUMN
             else:
                 local_annotations_column = next(k for k, v in col_mapping.items() if v == ANNOTATIONS_COLUMN)
-            if local_annotations_column not in t._schema.keys():
+            if local_annotations_column not in t._schema:
                 t.add_columns({local_annotations_column: pxt.JsonType(nullable=True)})
 
         resolved_col_mapping = cls.validate_columns(
@@ -591,9 +591,9 @@ class LabelStudioProject(Project):
             if media_import_method != 'url':
                 raise excs.Error("`s3_configuration` is only valid when `media_import_method == 'url'`")
             s3_configuration = copy.copy(s3_configuration)
-            if not 'bucket' in s3_configuration:
+            if 'bucket' not in s3_configuration:
                 raise excs.Error('`s3_configuration` must contain a `bucket` field')
-            if not 'title' in s3_configuration:
+            if 'title' not in s3_configuration:
                 s3_configuration['title'] = 'Pixeltable-S3-Import-Storage'
             if (
                 'aws_access_key_id' not in s3_configuration
@@ -633,7 +633,8 @@ class LabelStudioProject(Project):
                         raise excs.Error(
                             '`media_import_method` is set to `file`, but your Label Studio server is not configured '
                             'for local file storage.\nPlease set the `LABEL_STUDIO_LOCAL_FILES_SERVING_ENABLED` '
-                            'environment variable to `true` in the environment where your Label Studio server is running.'
+                            'environment variable to `true` in the environment where your Label Studio server '
+                            'is running.'
                         ) from exc
                 raise  # Handle any other exception type normally
 
@@ -663,7 +664,7 @@ class _LabelStudioConfig:
     rectangle_labels: dict[str, _RectangleLabel]
 
     def validate(self) -> None:
-        data_key_names = set(key.name for key in self.data_keys.values() if key.name is not None)
+        data_key_names = {key.name for key in self.data_keys.values() if key.name is not None}
         for name, rl in self.rectangle_labels.items():
             if rl.to_name not in data_key_names:
                 raise excs.Error(
@@ -674,7 +675,7 @@ class _LabelStudioConfig:
     @property
     def export_columns(self) -> dict[str, pxt.ColumnType]:
         data_key_cols = {key_id: key_info.column_type for key_id, key_info in self.data_keys.items()}
-        rl_cols = {name: pxt.JsonType() for name in self.rectangle_labels.keys()}
+        rl_cols = {name: pxt.JsonType() for name in self.rectangle_labels}
         return {**data_key_cols, **rl_cols}
 
 

pixeltable/io/pandas.py

@@ -1,3 +1,4 @@
+import os
 from typing import Any, Optional, Union
 
 import numpy as np
@@ -52,12 +53,12 @@ def import_pandas(
 
 def import_csv(
     tbl_name: str,
-    filepath_or_buffer,
+    filepath_or_buffer: Union[str, os.PathLike],
     schema_overrides: Optional[dict[str, Any]] = None,
     primary_key: Optional[Union[str, list[str]]] = None,
     num_retained_versions: int = 10,
     comment: str = '',
-    **kwargs,
+    **kwargs: Any,
 ) -> pxt.Table:
     """
     Creates a new base table from a csv file. This is a convenience method and is equivalent
@@ -81,13 +82,13 @@ def import_csv(
 
 def import_excel(
     tbl_name: str,
-    io,
-    *args,
+    io: Union[str, os.PathLike],
+    *,
     schema_overrides: Optional[dict[str, Any]] = None,
     primary_key: Optional[Union[str, list[str]]] = None,
     num_retained_versions: int = 10,
     comment: str = '',
-    **kwargs,
+    **kwargs: Any,
 ) -> pxt.Table:
     """
     Creates a new base table from an Excel (.xlsx) file. This is a convenience method and is

pixeltable/iterators/__init__.py

@@ -12,5 +12,5 @@ __removed_symbols = {'base', 'document', 'video'}
 __all__ = sorted(__default_dir - __removed_symbols)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/metadata/__init__.py

@@ -24,9 +24,11 @@ def create_system_info(engine: sql.engine.Engine) -> None:
     system_md = SystemInfoMd(schema_version=VERSION)
     record = SystemInfo(md=dataclasses.asdict(system_md))
     with orm.Session(engine, future=True) as session:
-        session.add(record)
-        session.flush()
-        session.commit()
+        # Write system metadata only once for idempotency
+        if session.query(SystemInfo).count() == 0:
+            session.add(record)
+            session.flush()
+            session.commit()
 
 
 # conversion functions for upgrading the metadata schema from one version to the following

pixeltable/metadata/converters/convert_24.py

@@ -19,11 +19,11 @@ def __substitute_md(k: Optional[str], v: Any) -> Optional[tuple[Optional[str], A
         isinstance(v, dict)
         and '_classpath' in v
         and v['_classpath']
-        in {
+        in (
             'pixeltable.func.callable_function.CallableFunction',
             'pixeltable.func.aggregate_function.AggregateFunction',
             'pixeltable.func.expr_template_function.ExprTemplateFunction',
-        }
+        )
     ):
         if 'path' in v:
             assert 'signature' not in v
@@ -50,6 +50,6 @@ def __substitute_path(path: str) -> str:
     # versions, it's necessary to resolve the function symbol to get the signature. The following
     # adjustment is necessary for function names that are stored in db artifacts of version < 25, but
     # have changed in some version > 25.
-    if path in {'pixeltable.functions.huggingface.clip_text', 'pixeltable.functions.huggingface.clip_image'}:
+    if path in ('pixeltable.functions.huggingface.clip_text', 'pixeltable.functions.huggingface.clip_image'):
         return 'pixeltable.functions.huggingface.clip'
     return path

pixeltable/metadata/converters/convert_25.py

@@ -13,7 +13,7 @@ def _(engine: sql.engine.Engine) -> None:
 
 def __substitute_md(k: Optional[str], v: Any) -> Optional[tuple[Optional[str], Any]]:
     if k == 'path' and (
-        v in {'pixeltable.functions.huggingface.clip_text', 'pixeltable.functions.huggingface.clip_image'}
+        v in ('pixeltable.functions.huggingface.clip_text', 'pixeltable.functions.huggingface.clip_image')
     ):
         return 'path', 'pixeltable.functions.huggingface.clip'
     return None

pixeltable/metadata/converters/convert_29.py

@@ -80,7 +80,7 @@ def __substitute_md(k: Optional[str], v: Any) -> Optional[tuple[Optional[str], A
                         rolled_kwargs = kwargs.pop(param['name'])
 
         if rolled_args is not None:
-            assert rolled_args['_classname'] in {'InlineArray', 'InlineList'}
+            assert rolled_args['_classname'] in ('InlineArray', 'InlineList')
             new_args.extend(rolled_args['components'])
         if rolled_kwargs is not None:
             assert rolled_kwargs['_classname'] == 'InlineDict'

pixeltable/store.py

@@ -99,9 +99,9 @@ class StoreBase:
 
         # v_min/v_max indices: speeds up base table scans needed to propagate a base table insert or delete
         idx_name = f'vmin_idx_{self.tbl_version.id.hex}'
-        idxs.append(sql.Index(idx_name, self.v_min_col, postgresql_using='brin'))
+        idxs.append(sql.Index(idx_name, self.v_min_col, postgresql_using=Env.get().dbms.version_index_type))
         idx_name = f'vmax_idx_{self.tbl_version.id.hex}'
-        idxs.append(sql.Index(idx_name, self.v_max_col, postgresql_using='brin'))
+        idxs.append(sql.Index(idx_name, self.v_max_col, postgresql_using=Env.get().dbms.version_index_type))
 
         self.sa_tbl = sql.Table(self._storage_name(), self.sa_md, *all_cols, *idxs)
 

pixeltable/type_system.py

@@ -512,7 +512,7 @@ class StringType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.STRING, nullable=nullable)
 
-    def has_supertype(self):
+    def has_supertype(self) -> bool:
         return not self.nullable
 
     @classmethod
@@ -602,7 +602,7 @@ class TimestampType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.TIMESTAMP, nullable=nullable)
 
-    def has_supertype(self):
+    def has_supertype(self) -> bool:
         return not self.nullable
 
     @classmethod
@@ -768,7 +768,7 @@ class JsonType(ColumnType):
         a_type = a.get('type')
         b_type = b.get('type')
 
-        if a_type in {'string', 'integer', 'number', 'boolean', 'object', 'array'} and a_type == b_type:
+        if a_type in ('string', 'integer', 'number', 'boolean', 'object', 'array') and a_type == b_type:
             # a and b both have the same type designation, but are not identical. This can happen if
             # (for example) they have validators or other attributes that differ. In this case, we
             # generalize to {'type': t}, where t is their shared type, with no other qualifications.
@@ -1170,6 +1170,20 @@ class DocumentType(ColumnType):
         XML = 3
         TXT = 4
 
+        @classmethod
+        def from_extension(cls, ext: str) -> Optional['DocumentType.DocumentFormat']:
+            if ext in ('.htm', '.html'):
+                return cls.HTML
+            if ext == '.md':
+                return cls.MD
+            if ext == '.pdf':
+                return cls.PDF
+            if ext == '.xml':
+                return cls.XML
+            if ext == '.txt':
+                return cls.TXT
+            return None
+
     def __init__(self, nullable: bool = False, doc_formats: Optional[str] = None):
         super().__init__(self.Type.DOCUMENT, nullable=nullable)
         self.doc_formats = doc_formats
@@ -1203,9 +1217,7 @@ class DocumentType(ColumnType):
         assert isinstance(val, str)
         from pixeltable.utils.documents import get_document_handle
 
-        dh = get_document_handle(val)
-        if dh is None:
-            raise excs.Error(f'Not a recognized document format: {val}')
+        _ = get_document_handle(val)
 
 
 T = typing.TypeVar('T')
@@ -1240,7 +1252,7 @@ class _PxtType:
     `ColumnType`.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         raise TypeError(f'Type `{type(self)}` cannot be instantiated.')
 
     @classmethod

pixeltable/utils/console_output.py

@@ -1,4 +1,5 @@
 import logging
+from typing import TextIO
 
 
 def map_level(verbosity: int) -> int:
@@ -22,10 +23,10 @@ def map_level(verbosity: int) -> int:
 
 
 class ConsoleOutputHandler(logging.StreamHandler):
-    def __init__(self, stream):
+    def __init__(self, stream: TextIO):
         super().__init__(stream)
 
-    def emit(self, record):
+    def emit(self, record: logging.LogRecord) -> None:
         if record.msg.endswith('\n'):
             self.stream.write(record.msg)
         else: