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

pixeltable/utils/coroutine.py

@@ -7,8 +7,8 @@ T = TypeVar('T')
 
 
 # TODO This is a temporary hack to be able to run async UDFs in contexts that are not properly handled by the existing
-#   scheduler logic (e.g., inside the eval loop of a JsonMapper). Once the scheduler is fully general, it can be
-#   removed.
+#   scheduler logic (e.g., as an embedding function as part of a similarity lookup). Once the scheduler is fully
+#   general, it can be removed.
 
 
 def run_coroutine_synchronously(coroutine: Coroutine[Any, Any, T], timeout: float = 30) -> T:
@@ -16,7 +16,7 @@ def run_coroutine_synchronously(coroutine: Coroutine[Any, Any, T], timeout: floa
     Runs the given coroutine synchronously, even if called in the context of a running event loop.
     """
 
-    def run_in_new_loop():
+    def run_in_new_loop() -> T:
         new_loop = asyncio.new_event_loop()
         asyncio.set_event_loop(new_loop)
         try:

pixeltable/utils/dbms.py

@@ -0,0 +1,66 @@
+import abc
+
+from sqlalchemy import URL
+
+
+class Dbms(abc.ABC):
+    """
+    Provides abstractions for utilities to interact with a database system.
+    """
+
+    name: str
+    transaction_isolation_level: str
+    version_index_type: str
+    db_url: URL
+
+    def __init__(self, name: str, transaction_isolation_level: str, version_index_type: str, db_url: URL) -> None:
+        self.name = name
+        self.transaction_isolation_level = transaction_isolation_level
+        self.version_index_type = version_index_type
+        self.db_url = db_url
+
+    @abc.abstractmethod
+    def drop_db_stmt(self, database: str) -> str: ...
+
+    @abc.abstractmethod
+    def create_db_stmt(self, database: str) -> str: ...
+
+    @abc.abstractmethod
+    def default_system_db_url(self) -> str: ...
+
+
+class PostgresqlDbms(Dbms):
+    """
+    Implements utilities to interact with Postgres database.
+    """
+
+    def __init__(self, db_url: URL):
+        super().__init__('postgresql', 'REPEATABLE READ', 'brin', db_url)
+
+    def drop_db_stmt(self, database: str) -> str:
+        return f'DROP DATABASE {database}'
+
+    def create_db_stmt(self, database: str) -> str:
+        return f"CREATE DATABASE {database} ENCODING 'utf-8' LC_COLLATE 'C' LC_CTYPE 'C' TEMPLATE template0"
+
+    def default_system_db_url(self) -> str:
+        a = self.db_url.set(database='postgres').render_as_string(hide_password=False)
+        return a
+
+
+class CockroachDbms(Dbms):
+    """
+    Implements utilities to interact with CockroachDb database.
+    """
+
+    def __init__(self, db_url: URL):
+        super().__init__('cockroachdb', 'SERIALIZABLE', 'btree', db_url)
+
+    def drop_db_stmt(self, database: str) -> str:
+        return f'DROP DATABASE {database} CASCADE'
+
+    def create_db_stmt(self, database: str) -> str:
+        return f"CREATE DATABASE {database} TEMPLATE template0 ENCODING 'utf-8' LC_COLLATE 'C' LC_CTYPE 'C'"
+
+    def default_system_db_url(self) -> str:
+        return self.db_url.set(database='defaultdb').render_as_string(hide_password=False)

pixeltable/utils/documents.py

@@ -1,11 +1,12 @@
 import dataclasses
+import os
 from typing import Optional
 
 import bs4
 import fitz  # type: ignore[import-untyped]
 import puremagic
 
-import pixeltable.type_system as ts
+from pixeltable import exceptions as excs, type_system as ts
 from pixeltable.env import Env
 
 
@@ -18,85 +19,78 @@ class DocumentHandle:
     txt_doc: Optional[str] = None
 
 
-def get_document_handle(path: str) -> Optional[DocumentHandle]:
-    doc_format = puremagic.from_file(path)
+def get_document_handle(path: str) -> DocumentHandle:
+    _, extension = os.path.splitext(path)
+    handle = get_handle_by_extension(path, extension)
+    if handle is not None:
+        return handle
 
-    if doc_format == '.pdf':
-        pdf_doc = get_pdf_handle(path)
-        if pdf_doc is not None:
-            return DocumentHandle(format=ts.DocumentType.DocumentFormat.PDF, pdf_doc=pdf_doc)
+    # if no extension, use puremagic to determine the type
+    extension = puremagic.from_file(path)
+    handle = get_handle_by_extension(path, extension)
+    if handle is not None:
+        return handle
 
-    if doc_format == '.html':
-        bs_doc = get_html_handle(path)
-        if bs_doc is not None:
-            return DocumentHandle(format=ts.DocumentType.DocumentFormat.HTML, bs_doc=bs_doc)
+    raise excs.Error(f'Unrecognized document format: {path}')
 
-    if doc_format == '.md':
-        md_ast = get_markdown_handle(path)
-        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)
+def get_handle_by_extension(path: str, extension: str) -> Optional[DocumentHandle]:
+    doc_format = ts.DocumentType.DocumentFormat.from_extension(extension)
 
-    if doc_format == '.txt':
-        txt_doc = get_txt(path)
-        if txt_doc is not None:
-            return DocumentHandle(format=ts.DocumentType.DocumentFormat.TXT, txt_doc=txt_doc)
+    try:
+        if doc_format == ts.DocumentType.DocumentFormat.HTML:
+            return DocumentHandle(doc_format, bs_doc=get_html_handle(path))
+        if doc_format == ts.DocumentType.DocumentFormat.MD:
+            return DocumentHandle(doc_format, md_ast=get_markdown_handle(path))
+        if doc_format == ts.DocumentType.DocumentFormat.PDF:
+            return DocumentHandle(doc_format, pdf_doc=get_pdf_handle(path))
+        if doc_format == ts.DocumentType.DocumentFormat.XML:
+            return DocumentHandle(doc_format, bs_doc=get_xml_handle(path))
+        if doc_format == ts.DocumentType.DocumentFormat.TXT:
+            return DocumentHandle(doc_format, txt_doc=get_txt(path))
+    except Exception as exc:
+        raise excs.Error(f'An error occurred processing a {doc_format} document: {path}') from exc
 
     return None
 
 
-def get_pdf_handle(path: str) -> Optional[fitz.Document]:
-    try:
-        doc = fitz.open(path)
-        # check pdf (bc it will work for images)
-        if not doc.is_pdf:
-            return None
-        # try to read one page
-        next(page for page in doc)
-        return doc
-    except Exception:
-        return None
-
-
-def get_html_handle(path: str) -> Optional[bs4.BeautifulSoup]:
-    try:
-        with open(path, 'r', encoding='utf8') as fp:
-            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
+def get_html_handle(path: str) -> bs4.BeautifulSoup:
+    with open(path, 'r', encoding='utf8') as fp:
+        doc = bs4.BeautifulSoup(fp, 'lxml')
+    if doc.find() is None:
+        raise excs.Error(f'Not a valid HTML document: {path}')
+    return doc
 
 
-def get_markdown_handle(path: str) -> Optional[dict]:
+def get_markdown_handle(path: str) -> dict:
     Env.get().require_package('mistune', [3, 0])
     import mistune
 
-    try:
-        with open(path, encoding='utf8') as file:
-            text = file.read()
-        md_ast = mistune.create_markdown(renderer=None)
-        return md_ast(text)
-    except Exception:
-        return None
+    with open(path, encoding='utf8') as file:
+        text = file.read()
+    md_ast = mistune.create_markdown(renderer=None)
+    return md_ast(text)
 
 
-def get_txt(path: str) -> Optional[str]:
-    try:
-        with open(path, 'r', encoding='utf-8') as fp:
-            doc = fp.read()
-        return doc or None  # replace '' with None
-    except Exception:
-        return None
+def get_pdf_handle(path: str) -> fitz.Document:
+    doc = fitz.open(path)
+    # check pdf (bc it will work for images)
+    if not doc.is_pdf:
+        raise excs.Error(f'Not a valid PDF document: {path}')
+    # try to read one page
+    next(page for page in doc)
+    return doc
+
+
+def get_xml_handle(path: str) -> bs4.BeautifulSoup:
+    with open(path, 'r', encoding='utf8') as fp:
+        doc = bs4.BeautifulSoup(fp, 'xml')
+    if doc.find() is None:
+        raise excs.Error(f'Not a valid XML document: {path}')
+    return doc
+
+
+def get_txt(path: str) -> str:
+    with open(path, 'r', encoding='utf-8') as fp:
+        doc = fp.read()
+    return doc

pixeltable/utils/exception_handler.py

@@ -0,0 +1,59 @@
+import logging
+import sys
+from typing import Any, Callable, Optional, TypeVar
+
+R = TypeVar('R')
+
+
+def _is_in_exception() -> bool:
+    """
+    Check if code is currently executing within an exception context.
+    """
+    current_exception = sys.exc_info()[1]
+    return current_exception is not None
+
+
+def run_cleanup_on_exception(cleanup_func: Callable[..., R], *args: Any, **kwargs: Any) -> Optional[R]:
+    """
+    Runs cleanup only when running in exception context.
+
+    The function `run_cleanup_on_exception()` should be used to clean up resources when an operation fails.
+    This is typically done using a try, except, and finally block, with the resource cleanup logic placed within
+    the except block. However, this pattern may not handle KeyboardInterrupt exceptions.
+    To ensure that resources are always cleaned up at least once when an exception or KeyboardInterrupt occurs,
+    create an idempotent function for cleaning up resources and pass it to the `run_cleanup_on_exception()` function
+    from the finally block.
+    """
+    if _is_in_exception():
+        return run_cleanup(cleanup_func, *args, raise_error=False, **kwargs)
+    return None
+
+
+def run_cleanup(cleanup_func: Callable[..., R], *args: Any, raise_error: bool = True, **kwargs: Any) -> Optional[R]:
+    """
+    Runs a cleanup function. If interrupted, retry cleanup.
+    The `run_cleanup()` function ensures that the `cleanup_func()` function executes at least once.
+    If the `cleanup_func()` is interrupted during execution, it will be retried.
+
+    Args:
+        cleanup_func: an idempotent function
+        raise_error: raise an exception if an error occurs during cleanup.
+    """
+    try:
+        logging.debug(f'Running cleanup function: {cleanup_func.__name__!r}')
+        return cleanup_func(*args, **kwargs)
+    except KeyboardInterrupt as interrupt:
+        # Save original exception and re-attempt cleanup
+        original_exception = interrupt
+        logging.debug(f'Cleanup {cleanup_func.__name__!r} interrupted, retrying')
+        try:
+            return cleanup_func(*args, **kwargs)
+        except Exception as e:
+            # Suppress this exception
+            logging.error(f'Cleanup {cleanup_func.__name__!r} failed with exception {e}')
+        raise KeyboardInterrupt from original_exception
+    except Exception as e:
+        logging.error(f'Cleanup {cleanup_func.__name__!r} failed with exception {e}')
+        if raise_error:
+            raise e
+    return None

pixeltable/utils/filecache.py

@@ -102,7 +102,7 @@ class FileCache:
     def init(cls) -> None:
         cls.__instance = cls()
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.cache = OrderedDict()
         self.total_size = 0
         self.capacity_bytes = int(Env.get()._file_cache_size_g * (1 << 30))

pixeltable/utils/http_server.py

@@ -3,6 +3,7 @@ import http.server
 import logging
 import pathlib
 import urllib
+from typing import Any
 
 _logger = logging.getLogger('pixeltable.http.server')
 
@@ -38,7 +39,7 @@ class AbsolutePathHandler(http.server.SimpleHTTPRequestHandler):
         path = pathlib.Path(urllib.request.url2pathname(path))
         return str(path)
 
-    def log_message(self, format, *args) -> None:
+    def log_message(self, format: str, *args: Any) -> None:
         """override logging to stderr in http.server.BaseHTTPRequestHandler"""
         message = format % args
         _logger.info(message.translate(self._control_char_table))  # type: ignore[attr-defined]
@@ -47,7 +48,7 @@ class AbsolutePathHandler(http.server.SimpleHTTPRequestHandler):
 class LoggingHTTPServer(http.server.ThreadingHTTPServer):
     """Avoids polluting stdout and stderr"""
 
-    def handle_error(self, request, client_address) -> None:
+    def handle_error(self, request, client_address) -> None:  # type: ignore[no-untyped-def]
         """override socketserver.TCPServer.handle_error which prints directly to sys.stderr"""
         import traceback
 

pixeltable/utils/pytorch.py

@@ -32,7 +32,7 @@ class PixeltablePytorchDataset(torch.utils.data.IterableDataset):
 
         self.path = path
         self.image_format = image_format
-        assert image_format in {'np', 'pt'}
+        assert image_format in ('np', 'pt')
         column_type_path = path / '.pixeltable.column_types.json'
         assert column_type_path.exists(), f'missing {column_type_path}'
         with column_type_path.open() as f:

pixeltable/utils/sql.py

@@ -4,7 +4,7 @@ import sqlalchemy as sql
 from sqlalchemy.dialects import postgresql
 
 
-def log_stmt(logger: logging.Logger, stmt) -> None:
+def log_stmt(logger: logging.Logger, stmt: sql.sql.ClauseElement) -> None:
     logger.debug(f'executing {stmt.compile(dialect=postgresql.dialect())}')