pixeltable 0.2.20__py3-none-any.whl → 0.2.22__py3-none-any.whl

pixeltable/type_system.py

@@ -3,20 +3,22 @@ from __future__ import annotations
 import abc
 import datetime
 import enum
+import io
 import json
 import typing
 import urllib.parse
 import urllib.request
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, Iterable, Mapping, Optional, Sequence, Union
 
+import PIL.Image
 import av  # type: ignore
 import numpy as np
-import PIL.Image
 import sqlalchemy as sql
+from typing import _GenericAlias  # type: ignore[attr-defined]
+from typing_extensions import _AnnotatedAlias
 
 import pixeltable.exceptions as excs
-from pixeltable.env import Env
 
 
 class ColumnType:
@@ -41,7 +43,7 @@ class ColumnType:
         def supertype(
                 cls, type1: 'ColumnType.Type', type2: 'ColumnType.Type',
                 # we need to pass this in because we can't easily append it as a class member
-                common_supertypes: Dict[Tuple['ColumnType.Type', 'ColumnType.Type'], 'ColumnType.Type']
+                common_supertypes: dict[tuple['ColumnType.Type', 'ColumnType.Type'], 'ColumnType.Type']
         ) -> Optional['ColumnType.Type']:
             if type1 == type2:
                 return type1
@@ -103,16 +105,16 @@ class ColumnType:
             return self.__class__(nullable=nullable)  # type: ignore[call-arg]
 
     @classmethod
-    def serialize_list(cls, type_list: List[ColumnType]) -> str:
+    def serialize_list(cls, type_list: list[ColumnType]) -> str:
         return json.dumps([t.as_dict() for t in type_list])
 
-    def as_dict(self) -> Dict:
+    def as_dict(self) -> dict:
         return {
             '_classname': self.__class__.__name__,
             **self._as_dict(),
         }
 
-    def _as_dict(self) -> Dict:
+    def _as_dict(self) -> dict:
         return {'nullable': self.nullable}
 
     @classmethod
@@ -121,18 +123,18 @@ class ColumnType:
         return cls.from_dict(type_dict)
 
     @classmethod
-    def deserialize_list(cls, type_list_str: str) -> List[ColumnType]:
+    def deserialize_list(cls, type_list_str: str) -> list[ColumnType]:
         type_dict_list = json.loads(type_list_str)
         return [cls.from_dict(type_dict) for type_dict in type_dict_list]
 
     @classmethod
-    def from_dict(cls, type_dict: Dict) -> ColumnType:
+    def from_dict(cls, type_dict: dict) -> ColumnType:
         assert '_classname' in type_dict
         type_class = globals()[type_dict['_classname']]
         return type_class._from_dict(type_dict)
 
     @classmethod
-    def _from_dict(cls, d: Dict) -> ColumnType:
+    def _from_dict(cls, d: dict) -> ColumnType:
         """
         Default implementation: simply invoke c'tor
         """
@@ -164,11 +166,28 @@ class ColumnType:
             return DocumentType()
 
     def __str__(self) -> str:
-        return self._type.name.lower()
+        return self._to_str(as_schema=False)
+
+    def _to_str(self, as_schema: bool) -> str:
+        base_str = self._to_base_str()
+        if as_schema:
+            return base_str if self.nullable else f'Required[{base_str}]'
+        else:
+            return f'Optional[{base_str}]' if self.nullable else base_str
+
+    def _to_base_str(self) -> str:
+        """
+        String representation of this type, disregarding nullability. Default implementation is to camel-case
+        the type name; subclasses may override.
+        """
+        return self._type.name[0] + self._type.name[1:].lower()
 
     def __eq__(self, other: object) -> bool:
         return isinstance(other, ColumnType) and self.matches(other) and self.nullable == other.nullable
 
+    def __hash__(self) -> int:
+        return hash((self._type, self.nullable))
+
     def is_supertype_of(self, other: ColumnType, ignore_nullable: bool = False) -> bool:
         if ignore_nullable:
             supertype = self.supertype(other)
@@ -253,39 +272,63 @@ class ColumnType:
         return inferred_type
 
     @classmethod
-    def from_python_type(cls, t: type) -> Optional[ColumnType]:
+    def from_python_type(cls, t: Union[type, _GenericAlias], nullable_default: bool = False) -> Optional[ColumnType]:
         if typing.get_origin(t) is typing.Union:
             union_args = typing.get_args(t)
-            if union_args[1] is type(None):
-                # `t` is a type of the form Optional[T] (equivalently, Union[T, None]).
+            if len(union_args) == 2 and type(None) in union_args:
+                # `t` is a type of the form Optional[T] (equivalently, Union[T, None] or Union[None, T]).
                 # We treat it as the underlying type but with nullable=True.
-                underlying = cls.from_python_type(union_args[0])
+                underlying_py_type = union_args[0] if union_args[1] is type(None) else union_args[1]
+                underlying = cls.from_python_type(underlying_py_type)
                 if underlying is not None:
-                    underlying._nullable = True
-                    return underlying
+                    return underlying.copy(nullable=True)
+        elif typing.get_origin(t) is typing.Annotated:
+            annotated_args = typing.get_args(t)
+            origin = annotated_args[0]
+            parameters = annotated_args[1]
+            if isinstance(parameters, ColumnType):
+                return parameters.copy(nullable=nullable_default)
+        elif typing.get_origin(t) is Required:
+            required_args = typing.get_args(t)
+            assert len(required_args) == 1
+            return cls.from_python_type(required_args[0], nullable_default=False)
         else:
             # Discard type parameters to ensure that parameterized types such as `list[T]`
             # are correctly mapped to Pixeltable types.
-            base = typing.get_origin(t)
-            if base is None:
-                # No type parameters; the base type is just `t` itself
-                base = t
-            if base is str:
-                return StringType()
-            if base is int:
-                return IntType()
-            if base is float:
-                return FloatType()
-            if base is bool:
-                return BoolType()
-            if base is datetime.datetime:
-                return TimestampType()
-            if issubclass(base, Sequence) or issubclass(base, Mapping):
-                return JsonType()
-            if issubclass(base, PIL.Image.Image):
-                return ImageType()
+            origin = typing.get_origin(t)
+            if origin is None:
+                # No type parameters; the origin type is just `t` itself
+                origin = t
+            if issubclass(origin, _PxtType):
+                return origin.as_col_type(nullable=nullable_default)
+            if origin is str:
+                return StringType(nullable=nullable_default)
+            if origin is int:
+                return IntType(nullable=nullable_default)
+            if origin is float:
+                return FloatType(nullable=nullable_default)
+            if origin is bool:
+                return BoolType(nullable=nullable_default)
+            if origin is datetime.datetime:
+                return TimestampType(nullable=nullable_default)
+            if origin is PIL.Image.Image:
+                return ImageType(nullable=nullable_default)
+            if issubclass(origin, Sequence) or issubclass(origin, Mapping):
+                return JsonType(nullable=nullable_default)
         return None
 
+    @classmethod
+    def normalize_type(cls, t: Union[ColumnType, type, _AnnotatedAlias], nullable_default: bool = False) -> ColumnType:
+        """
+        Convert any type recognizable by Pixeltable to its corresponding ColumnType.
+        """
+        if isinstance(t, ColumnType):
+            return t
+        col_type = cls.from_python_type(t, nullable_default)
+        if col_type is None:
+            raise excs.Error(f'Unknown type: {t}')
+        return col_type
+
     def validate_literal(self, val: Any) -> None:
         """Raise TypeError if val is not a valid literal for this type"""
         if val is None:
@@ -491,7 +534,7 @@ class TimestampType(ColumnType):
 
 class JsonType(ColumnType):
     # TODO: type_spec also needs to be able to express lists
-    def __init__(self, type_spec: Optional[Dict[str, ColumnType]] = None, nullable: bool = False):
+    def __init__(self, type_spec: Optional[dict[str, ColumnType]] = None, nullable: bool = False):
         super().__init__(self.Type.JSON, nullable=nullable)
         self.type_spec = type_spec
 
@@ -526,7 +569,7 @@ class JsonType(ColumnType):
                 type_spec[other_field_name] = field_type
         return JsonType(type_spec, nullable=(self.nullable or other.nullable))
 
-    def _as_dict(self) -> Dict:
+    def _as_dict(self) -> dict:
         result = super()._as_dict()
         if self.type_spec is not None:
             type_spec_dict = {field_name: field_type.serialize() for field_name, field_type in self.type_spec.items()}
@@ -534,7 +577,7 @@ class JsonType(ColumnType):
         return result
 
     @classmethod
-    def _from_dict(cls, d: Dict) -> ColumnType:
+    def _from_dict(cls, d: dict) -> ColumnType:
         type_spec = None
         if 'type_spec' in d:
             type_spec = {
@@ -590,6 +633,9 @@ class ArrayType(ColumnType):
     def matches(self, other: ColumnType) -> bool:
         return isinstance(other, ArrayType) and self.shape == other.shape and self.dtype == other.dtype
 
+    def __hash__(self) -> int:
+        return hash((self._type, self.nullable, self.shape, self.dtype))
+
     def supertype(self, other: ColumnType) -> Optional[ArrayType]:
         if not isinstance(other, ArrayType):
             return None
@@ -601,16 +647,16 @@ class ArrayType(ColumnType):
         shape = [n1 if n1 == n2 else None for n1, n2 in zip(self.shape, other.shape)]
         return ArrayType(tuple(shape), self.make_type(base_type), nullable=(self.nullable or other.nullable))
 
-    def _as_dict(self) -> Dict:
+    def _as_dict(self) -> dict:
         result = super()._as_dict()
         result.update(shape=list(self.shape), dtype=self.dtype.value)
         return result
 
-    def __str__(self) -> str:
-        return f'{self._type.name.lower()}({self.shape}, dtype={self.dtype.name})'
+    def _to_base_str(self) -> str:
+        return f'Array[{self.shape}, {self.pxt_dtype}]'
 
     @classmethod
-    def _from_dict(cls, d: Dict) -> ColumnType:
+    def _from_dict(cls, d: dict) -> ColumnType:
         assert 'shape' in d
         assert 'dtype' in d
         shape = tuple(d['shape'])
@@ -681,7 +727,7 @@ class ArrayType(ColumnType):
 
 class ImageType(ColumnType):
     def __init__(
-            self, width: Optional[int] = None, height: Optional[int] = None, size: Optional[Tuple[int, int]] = None,
+            self, width: Optional[int] = None, height: Optional[int] = None, size: Optional[tuple[int, int]] = None,
             mode: Optional[str] = None, nullable: bool = False
     ):
         """
@@ -701,23 +747,17 @@ class ImageType(ColumnType):
     def copy(self, nullable: bool) -> ColumnType:
         return ImageType(self.width, self.height, mode=self.mode, nullable=nullable)
 
-    def __str__(self) -> str:
-        if self.width is not None or self.height is not None or self.mode is not None:
+    def _to_base_str(self) -> str:
+        params = []
+        if self.width is not None or self.height is not None:
+            params.append(f'({self.width}, {self.height})')
+        if self.mode is not None:
+            params.append(repr(self.mode))
+        if len(params) == 0:
             params_str = ''
-            if self.width is not None:
-                params_str = f'width={self.width}'
-            if self.height is not None:
-                if len(params_str) > 0:
-                    params_str += ', '
-                params_str += f'height={self.height}'
-            if self.mode is not None:
-                if len(params_str) > 0:
-                    params_str += ', '
-                params_str += f'mode={self.mode}'
-            params_str = f'({params_str})'
         else:
-            params_str = ''
-        return f'{self._type.name.lower()}{params_str}'
+            params_str = f'[{", ".join(params)}]'
+        return f'Image{params_str}'
 
     def matches(self, other: ColumnType) -> bool:
         return (
@@ -727,6 +767,9 @@ class ImageType(ColumnType):
             and self.mode == other.mode
         )
 
+    def __hash__(self) -> int:
+        return hash((self._type, self.nullable, self.size, self.mode))
+
     def supertype(self, other: ColumnType) -> Optional[ImageType]:
         if not isinstance(other, ImageType):
             return None
@@ -736,18 +779,18 @@ class ImageType(ColumnType):
         return ImageType(width=width, height=height, mode=mode, nullable=(self.nullable or other.nullable))
 
     @property
-    def size(self) -> Optional[Tuple[int, int]]:
+    def size(self) -> Optional[tuple[int, int]]:
         if self.width is None or self.height is None:
             return None
         return (self.width, self.height)
 
-    def _as_dict(self) -> Dict:
+    def _as_dict(self) -> dict:
         result = super()._as_dict()
         result.update(width=self.width, height=self.height, mode=self.mode)
         return result
 
     @classmethod
-    def _from_dict(cls, d: Dict) -> ColumnType:
+    def _from_dict(cls, d: dict) -> ColumnType:
         assert 'width' in d
         assert 'height' in d
         assert 'mode' in d
@@ -756,6 +799,20 @@ class ImageType(ColumnType):
     def to_sa_type(self) -> sql.types.TypeEngine:
         return sql.String()
 
+    def _create_literal(self, val: Any) -> Any:
+        if isinstance(val, str) and val.startswith('data:'):
+            # try parsing this as a `data:` URL, and if successful, decode the image immediately
+            try:
+                with urllib.request.urlopen(val) as response:
+                    b = response.read()
+                img = PIL.Image.open(io.BytesIO(b))
+                img.load()
+                return img
+            except Exception as exc:
+                errormsg_val = val if len(val) < 50 else val[:50] + '...'
+                raise excs.Error(f'data URL could not be decoded into a valid image: {errormsg_val}') from exc
+        return val
+
     def _validate_literal(self, val: Any) -> None:
         if isinstance(val, PIL.Image.Image):
             return
@@ -834,6 +891,7 @@ class DocumentType(ColumnType):
         HTML = 0
         MD = 1
         PDF = 2
+        XML = 3
 
     def __init__(self, nullable: bool = False, doc_formats: Optional[str] = None):
         super().__init__(self.Type.DOCUMENT, nullable=nullable)
@@ -853,6 +911,9 @@ class DocumentType(ColumnType):
     def matches(self, other: ColumnType) -> bool:
         return isinstance(other, DocumentType) and self._doc_formats == other._doc_formats
 
+    def __hash__(self) -> int:
+        return hash((self._type, self.nullable, self._doc_formats))
+
     def to_sa_type(self) -> sql.types.TypeEngine:
         # stored as a file path
         return sql.String()
@@ -866,3 +927,141 @@ class DocumentType(ColumnType):
         dh = get_document_handle(val)
         if dh is None:
             raise excs.Error(f'Not a recognized document format: {val}')
+
+
+T = typing.TypeVar('T')
+
+
+class Required(typing.Generic[T]):
+    """
+    Marker class to indicate that a column is non-nullable in a schema definition. This has no meaning as a type hint,
+    and is intended only for schema declarations.
+    """
+    pass
+
+
+String = typing.Annotated[str, StringType(nullable=False)]
+Int = typing.Annotated[int, IntType(nullable=False)]
+Float = typing.Annotated[float, FloatType(nullable=False)]
+Bool = typing.Annotated[bool, BoolType(nullable=False)]
+Timestamp = typing.Annotated[datetime.datetime, TimestampType(nullable=False)]
+
+
+class _PxtType:
+    """
+    Base class for the Pixeltable type-hint family. Subclasses of this class are meant to be used as type hints, both
+    in schema definitions and in UDF signatures. Whereas `ColumnType`s are instantiable and carry semantic information
+    about the Pixeltable type system, `_PxtType` subclasses are purely for convenience: they are not instantiable and
+    must be resolved to a `ColumnType` (by calling `ColumnType.from_python_type()`) in order to do anything meaningful
+    with them.
+
+    `_PxtType` subclasses can be specialized (as type hints) with type parameters; for example:
+    `Image[(300, 300), 'RGB']`. The specialized forms resolve to `typing.Annotated` instances whose annotation is a
+    `ColumnType`.
+    """
+    def __init__(self):
+        raise TypeError(f'Type `{type(self)}` cannot be instantiated.')
+
+    @classmethod
+    def as_col_type(cls, nullable: bool) -> ColumnType:
+        raise NotImplementedError()
+
+
+class Json(_PxtType):
+    @classmethod
+    def as_col_type(cls, nullable: bool) -> ColumnType:
+        return JsonType(nullable=nullable)
+
+
+class Array(np.ndarray, _PxtType):
+    def __class_getitem__(cls, item: Any) -> _AnnotatedAlias:
+        """
+        `item` (the type subscript) must be a tuple with exactly two elements (in any order):
+        - A tuple of `Optional[int]`s, specifying the shape of the array
+        - A type, specifying the dtype of the array
+        Example: Array[(3, None, 2), float]
+        """
+        params = item if isinstance(item, tuple) else (item,)
+        shape: Optional[tuple] = None
+        dtype: Optional[ColumnType] = None
+        for param in params:
+            if isinstance(param, tuple):
+                if not all(n is None or (isinstance(n, int) and n >= 1) for n in param):
+                    raise TypeError(f'Invalid Array type parameter: {param}')
+                if shape is not None:
+                    raise TypeError(f'Duplicate Array type parameter: {param}')
+                shape = param
+            elif isinstance(param, type) or isinstance(param, _AnnotatedAlias):
+                if dtype is not None:
+                    raise TypeError(f'Duplicate Array type parameter: {param}')
+                dtype = ColumnType.from_python_type(param)
+            else:
+                raise TypeError(f'Invalid Array type parameter: {param}')
+        if shape is None:
+            raise TypeError('Array type is missing parameter: shape')
+        if dtype is None:
+            raise TypeError('Array type is missing parameter: dtype')
+        return typing.Annotated[np.ndarray, ArrayType(shape=shape, dtype=dtype, nullable=False)]
+
+    @classmethod
+    def as_col_type(cls, nullable: bool) -> ColumnType:
+        raise TypeError('Array type cannot be used without specifying shape and dtype')
+
+
+class Image(PIL.Image.Image, _PxtType):
+    def __class_getitem__(cls, item: Any) -> _AnnotatedAlias:
+        """
+        `item` (the type subscript) must be one of the following, or a tuple containing either or both in any order:
+        - A 2-tuple of `int`s, specifying the size of the image
+        - A string, specifying the mode of the image
+        Example: Image[(300, 300), 'RGB']
+        """
+        if isinstance(item, tuple) and all(n is None or isinstance(n, int) for n in item):
+            # It's a tuple of the form (width, height)
+            params = (item,)
+        elif isinstance(item, tuple):
+            # It's a compound tuple (multiple parameters)
+            params = item
+        else:
+            # Not a tuple (single arg)
+            params = (item,)
+        size: Optional[tuple] = None
+        mode: Optional[str] = None
+        for param in params:
+            if isinstance(param, tuple):
+                if len(param) != 2 or not isinstance(param[0], (int, type(None))) or not isinstance(param[1], (int, type(None))):
+                    raise TypeError(f'Invalid Image type parameter: {param}')
+                if size is not None:
+                    raise TypeError(f'Duplicate Image type parameter: {param}')
+                size = param
+            elif isinstance(param, str):
+                if param not in PIL.Image.MODES:
+                    raise TypeError(f'Invalid Image type parameter: {param!r}')
+                if mode is not None:
+                    raise TypeError(f'Duplicate Image type parameter: {param!r}')
+                mode = param
+            else:
+                raise TypeError(f'Invalid Image type parameter: {param}')
+        return typing.Annotated[PIL.Image.Image, ImageType(size=size, mode=mode, nullable=False)]
+
+    @classmethod
+    def as_col_type(cls, nullable: bool) -> ColumnType:
+        return ImageType(nullable=nullable)
+
+
+class Video(str, _PxtType):
+    @classmethod
+    def as_col_type(cls, nullable: bool) -> ColumnType:
+        return VideoType(nullable=nullable)
+
+
+class Audio(str, _PxtType):
+    @classmethod
+    def as_col_type(cls, nullable: bool) -> ColumnType:
+        return AudioType(nullable=nullable)
+
+
+class Document(str, _PxtType):
+    @classmethod
+    def as_col_type(cls, nullable: bool) -> ColumnType:
+        return DocumentType(nullable=nullable)

pixeltable/utils/arrow.py

@@ -1,13 +1,14 @@
 import logging
-from typing import Any, Dict, Iterable, Iterator, Optional
+from typing import Any, Iterator, Optional, Union
 
+import numpy as np
 import pyarrow as pa
 
 import pixeltable.type_system as ts
 
 _logger = logging.getLogger(__name__)
 
-_pa_to_pt: Dict[pa.DataType, ts.ColumnType] = {
+_pa_to_pt: dict[pa.DataType, ts.ColumnType] = {
     pa.string(): ts.StringType(nullable=True),
     pa.timestamp('us'): ts.TimestampType(nullable=True),
     pa.bool_(): ts.BoolType(nullable=True),
@@ -20,7 +21,7 @@ _pa_to_pt: Dict[pa.DataType, ts.ColumnType] = {
     pa.float32(): ts.FloatType(nullable=True),
 }
 
-_pt_to_pa: Dict[ts.ColumnType, pa.DataType] = {
+_pt_to_pa: dict[type[ts.ColumnType], pa.DataType] = {
     ts.StringType: pa.string(),
     ts.TimestampType: pa.timestamp('us'),  # postgres timestamp is microseconds
     ts.BoolType: pa.bool_(),
@@ -61,19 +62,19 @@ def to_arrow_type(pixeltable_type: ts.ColumnType) -> Optional[pa.DataType]:
         return None
 
 
-def to_pixeltable_schema(arrow_schema: pa.Schema) -> Dict[str, ts.ColumnType]:
+def to_pixeltable_schema(arrow_schema: pa.Schema) -> dict[str, ts.ColumnType]:
     return {field.name: to_pixeltable_type(field.type) for field in arrow_schema}
 
 
-def to_arrow_schema(pixeltable_schema: Dict[str, Any]) -> pa.Schema:
-    return pa.schema((name, to_arrow_type(typ)) for name, typ in pixeltable_schema.items())
+def to_arrow_schema(pixeltable_schema: dict[str, Any]) -> pa.Schema:
+    return pa.schema((name, to_arrow_type(typ)) for name, typ in pixeltable_schema.items())  # type: ignore[misc]
 
 
-def to_pydict(batch: pa.RecordBatch) -> Dict[str, Iterable[Any]]:
+def to_pydict(batch: pa.RecordBatch) -> dict[str, Union[list, np.ndarray]]:
     """Convert a RecordBatch to a dictionary of lists, unlike pa.lib.RecordBatch.to_pydict,
     this function will not convert numpy arrays to lists, and will preserve the original numpy dtype.
     """
-    out = {}
+    out: dict[str, Union[list, np.ndarray]] = {}
     for k, name in enumerate(batch.schema.names):
         col = batch.column(k)
         if isinstance(col.type, pa.FixedShapeTensorType):
@@ -86,7 +87,7 @@ def to_pydict(batch: pa.RecordBatch) -> Dict[str, Iterable[Any]]:
     return out
 
 
-def iter_tuples(batch: pa.RecordBatch) -> Iterator[Dict[str, Any]]:
+def iter_tuples(batch: pa.RecordBatch) -> Iterator[dict[str, Any]]:
     """Convert a RecordBatch to an iterator of dictionaries. also works with pa.Table and pa.RowGroup"""
     pydict = to_pydict(batch)
     assert len(pydict) > 0, 'empty record batch'

pixeltable/utils/coco.py

@@ -1,12 +1,12 @@
-from typing import List, Dict, Any, Set
-from pathlib import Path
 import json
+from pathlib import Path
+from typing import Any, Dict, List, Set
 
 import PIL
 
+import pixeltable as pxt
 import pixeltable.exceptions as excs
 
-
 format_msg = """
 
 Required format:
@@ -48,7 +48,7 @@ def _verify_input_dict(input_dict: Dict[str, Any]) -> None:
         if not isinstance(annotation['category'], (str, int)):
             raise excs.Error(f'Value for "category" is not a str or int: {annotation}{format_msg}')
 
-def write_coco_dataset(df: 'pixeltable.DataFrame', dest_path: Path) -> Path:
+def write_coco_dataset(df: pxt.DataFrame, dest_path: Path) -> Path:
     """Export a DataFrame result set as a COCO dataset in dest_path and return the path of the data.json file."""
     # TODO: validate schema
     if len(df._select_list_exprs) != 1 or not df._select_list_exprs[0].col_type.is_json_type():

pixeltable/utils/documents.py

@@ -2,7 +2,7 @@ import dataclasses
 from typing import Optional
 
 import bs4
-import fitz  # (pymupdf)
+import fitz  # type: ignore[import-untyped]
 import puremagic
 
 import pixeltable.type_system as ts
@@ -35,6 +35,11 @@ def get_document_handle(path: str) -> Optional[DocumentHandle]:
         if md_ast is not None:
             return DocumentHandle(format=ts.DocumentType.DocumentFormat.MD, md_ast=md_ast)
 
+    if doc_format == '.xml':
+        bs_doc = get_xml_handle(path)
+        if bs_doc is not None:
+            return DocumentHandle(format=ts.DocumentType.DocumentFormat.XML, bs_doc=bs_doc)
+
     return None
 
 
@@ -54,7 +59,16 @@ def get_pdf_handle(path: str) -> Optional[fitz.Document]:
 def get_html_handle(path: str) -> Optional[bs4.BeautifulSoup]:
     try:
         with open(path, 'r', encoding='utf8') as fp:
-            doc = bs4.BeautifulSoup(fp, 'html.parser')
+            doc = bs4.BeautifulSoup(fp, 'lxml')
+        return doc if doc.find() is not None else None
+    except Exception:
+        return None
+
+
+def get_xml_handle(path: str) -> Optional[bs4.BeautifulSoup]:
+    try:
+        with open(path, 'r', encoding='utf8') as fp:
+            doc = bs4.BeautifulSoup(fp, 'xml')
         return doc if doc.find() is not None else None
     except Exception:
         return None