pixeltable 0.4.0rc3__py3-none-any.whl → 0.4.20__py3-none-any.whl

pixeltable/utils/av.py

@@ -0,0 +1,298 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from fractions import Fraction
+from pathlib import Path
+from types import TracebackType
+from typing import Any, Iterator
+
+import av
+import av.stream
+import PIL.Image
+from typing_extensions import Self
+
+from pixeltable.env import Env
+
+# format -> (codec, extension)
+AUDIO_FORMATS: dict[str, tuple[str, str]] = {
+    'wav': ('pcm_s16le', 'wav'),
+    'mp3': ('libmp3lame', 'mp3'),
+    'flac': ('flac', 'flac'),
+    'mp4': ('aac', 'm4a'),
+}
+
+
+def get_metadata(path: str) -> dict:
+    with av.open(path) as container:
+        assert isinstance(container, av.container.InputContainer)
+        streams_info = [__get_stream_metadata(stream) for stream in container.streams]
+        result = {
+            'bit_exact': getattr(container, 'bit_exact', False),
+            'bit_rate': container.bit_rate,
+            'size': container.size,
+            'metadata': container.metadata,
+            'streams': streams_info,
+        }
+    return result
+
+
+def __get_stream_metadata(stream: av.stream.Stream) -> dict:
+    if stream.type not in ('audio', 'video'):
+        return {'type': stream.type}  # Currently unsupported
+
+    codec_context = stream.codec_context
+    codec_context_md: dict[str, Any] = {
+        'name': codec_context.name,
+        'codec_tag': codec_context.codec_tag.encode('unicode-escape').decode('utf-8'),
+        'profile': codec_context.profile,
+    }
+    metadata = {
+        'type': stream.type,
+        'duration': stream.duration,
+        'time_base': float(stream.time_base) if stream.time_base is not None else None,
+        'duration_seconds': float(stream.duration * stream.time_base)
+        if stream.duration is not None and stream.time_base is not None
+        else None,
+        'frames': stream.frames,
+        'metadata': stream.metadata,
+        'codec_context': codec_context_md,
+    }
+
+    if stream.type == 'audio':
+        # Additional metadata for audio
+        channels = getattr(stream.codec_context, 'channels', None)
+        codec_context_md['channels'] = int(channels) if channels is not None else None
+    else:
+        assert stream.type == 'video'
+        assert isinstance(stream, av.video.stream.VideoStream)
+        # Additional metadata for video
+        codec_context_md['pix_fmt'] = getattr(stream.codec_context, 'pix_fmt', None)
+        metadata.update(
+            **{
+                'width': stream.width,
+                'height': stream.height,
+                'frames': stream.frames,
+                'average_rate': float(stream.average_rate) if stream.average_rate is not None else None,
+                'base_rate': float(stream.base_rate) if stream.base_rate is not None else None,
+                'guessed_rate': float(stream.guessed_rate) if stream.guessed_rate is not None else None,
+            }
+        )
+
+    return metadata
+
+
+def get_video_duration(path: str) -> float | None:
+    """Return video duration in seconds."""
+    with av.open(path) as container:
+        video_stream = container.streams.video[0]
+        if video_stream is None:
+            return None
+        if video_stream.duration is not None:
+            return float(video_stream.duration * video_stream.time_base)
+
+        # if duration is not in the header, look for it in the last packet
+        last_pts: int | None = None
+        for packet in container.demux(video_stream):
+            if packet.pts is not None:
+                last_pts = packet.pts
+        if last_pts is not None:
+            return float(last_pts * video_stream.time_base)
+
+        return None
+
+
+def has_audio_stream(path: str) -> bool:
+    """Check if video has audio stream using PyAV."""
+    md = get_metadata(path)
+    return any(stream['type'] == 'audio' for stream in md['streams'])
+
+
+def ffmpeg_clip_cmd(
+    input_path: str,
+    output_path: str,
+    start_time: float,
+    duration: float | None = None,
+    fast: bool = True,
+    video_encoder: str | None = None,
+    video_encoder_args: dict[str, Any] | None = None,
+) -> list[str]:
+    cmd = ['ffmpeg']
+    if fast:
+        # fast: -ss before -i
+        cmd.extend(
+            [
+                '-ss',
+                str(start_time),
+                '-i',
+                input_path,
+                '-map',
+                '0',  # Copy all streams from input
+                '-c',
+                'copy',  # Stream copy (no re-encoding)
+            ]
+        )
+    else:
+        if video_encoder is None:
+            video_encoder = Env.get().default_video_encoder
+
+        # accurate: -ss after -i
+        cmd.extend(
+            [
+                '-i',
+                input_path,
+                '-ss',
+                str(start_time),
+                '-map',
+                '0',  # Copy all streams from input
+                '-c:a',
+                'copy',  # audio copy
+                '-c:s',
+                'copy',  # subtitle copy
+                '-c:v',
+                video_encoder,  # re-encode video
+            ]
+        )
+        if video_encoder_args is not None:
+            for k, v in video_encoder_args.items():
+                cmd.extend([f'-{k}', str(v)])
+
+    if duration is not None:
+        cmd.extend(['-t', str(duration)])
+    cmd.extend(['-loglevel', 'error', output_path])
+    return cmd
+
+
+def ffmpeg_segment_cmd(
+    input_path: str,
+    output_pattern: str,
+    segment_duration: float | None = None,
+    segment_times: list[float] | None = None,
+    video_encoder: str | None = None,
+    video_encoder_args: dict[str, Any] | None = None,
+) -> list[str]:
+    """Commandline for frame-accurate segmentation"""
+    assert (segment_duration is None) != (segment_times is None)
+    if video_encoder is None:
+        video_encoder = Env.get().default_video_encoder
+
+    cmd = [
+        'ffmpeg',
+        '-i',
+        input_path,
+        '-map',
+        '0',  # Copy all streams from input
+        '-c:a',
+        'copy',  # don't re-encode audio
+        '-c:v',
+        video_encoder,  # re-encode video
+    ]
+    if video_encoder_args is not None:
+        for k, v in video_encoder_args.items():
+            cmd.extend([f'-{k}', str(v)])
+    cmd.extend(['-f', 'segment'])
+
+    # -force_key_frames needs to precede -f segment
+    if segment_duration is not None:
+        cmd.extend(
+            [
+                '-force_key_frames',
+                f'expr:gte(t,n_forced*{segment_duration})',  # Force keyframe at each segment boundary
+                '-f',
+                'segment',
+                '-segment_time',
+                str(segment_duration),
+            ]
+        )
+    else:
+        assert segment_times is not None
+        times_str = ','.join([str(t) for t in segment_times])
+        cmd.extend(['-force_key_frames', times_str, '-f', 'segment', '-segment_times', times_str])
+
+    cmd.extend(
+        [
+            '-reset_timestamps',
+            '1',  # Reset timestamps for each segment
+            '-loglevel',
+            'error',  # Only show errors
+            output_pattern,
+        ]
+    )
+    return cmd
+
+
+class VideoFrames:
+    """
+    Context manager for iterating over video frames at a specified frame rate.
+
+    Args:
+        path: Path to the video file
+        fps: Number of frames to extract per second. If None or 0.0, extracts all frames.
+    """
+
+    path: Path
+    fps: float
+    container: av.container.input.InputContainer | None
+    video_framerate: Fraction | None
+    video_time_base: Fraction | None
+    video_start_time: int | None
+
+    @dataclass
+    class Item:
+        frame_idx: int
+        pts: int
+        dts: int
+        time: float
+        is_corrupt: bool
+        key_frame: bool
+        pict_type: int
+        interlaced_frame: bool
+        frame: PIL.Image.Image
+
+    def __init__(self, path: Path, fps: float | None = None) -> None:
+        self.path = path
+        self.fps = 0.0 if fps is None else fps
+        self.container = None
+        self.video_framerate = None
+        self.video_time_base = None
+        self.video_start_time = None
+
+    def __enter__(self) -> Self:
+        self.container = av.open(self.path)
+        stream = self.container.streams.video[0]
+        self.video_framerate = stream.average_rate
+        self.video_time_base = stream.time_base
+        self.video_start_time = stream.start_time or 0
+        return self
+
+    def __exit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
+    ) -> None:
+        # Clean up
+        if self.container:
+            self.container.close()
+
+    def __iter__(self) -> Iterator[Item]:
+        num_returned = 0
+        frame_idx = -1
+        while True:
+            try:
+                frame = next(self.container.decode(video=0))
+            except (StopIteration, EOFError):
+                return
+
+            frame_idx += 1
+            if self.fps == 0.0 or (num_returned <= frame.time * self.fps):
+                img = frame.to_image()
+                assert isinstance(img, PIL.Image.Image)
+                yield VideoFrames.Item(
+                    frame_idx=frame_idx,
+                    pts=frame.pts,
+                    dts=frame.dts,
+                    time=frame.time,
+                    is_corrupt=frame.is_corrupt,
+                    key_frame=frame.key_frame,
+                    pict_type=frame.pict_type,
+                    interlaced_frame=frame.interlaced_frame,
+                    frame=img,
+                )
+                num_returned += 1

pixeltable/utils/azure_store.py

@@ -0,0 +1,305 @@
+import logging
+import re
+import threading
+import uuid
+from pathlib import Path
+from typing import TYPE_CHECKING, Iterator
+
+from azure.core.exceptions import AzureError
+
+from pixeltable import env, exceptions as excs
+from pixeltable.config import Config
+from pixeltable.utils.object_stores import ObjectPath, ObjectStoreBase, StorageObjectAddress
+
+if TYPE_CHECKING:
+    from azure.storage.blob import BlobProperties, BlobServiceClient
+
+    from pixeltable.catalog import Column
+
+
+_logger = logging.getLogger('pixeltable')
+
+
+client_lock = threading.Lock()
+
+
+@env.register_client('azure_blob')
+def _() -> dict[str, 'BlobServiceClient']:
+    return {}
+
+
+class AzureBlobStore(ObjectStoreBase):
+    """Class to handle Azure Blob Storage operations."""
+
+    # TODO: This needs to be redesigned to use asyncio.
+
+    # URI of the Azure Blob Storage container
+    # Always ends with a slash
+    __base_uri: str
+
+    # Storage account name
+    __account_name: str
+
+    # Container name extracted from the URI
+    __container_name: str
+
+    # Prefix path within the container, either empty or ending with a slash
+    __prefix_name: str
+
+    # URI scheme (wasb, wasbs, abfs, abfss, https)
+    __scheme: str
+
+    soa: StorageObjectAddress
+
+    def __init__(self, soa: StorageObjectAddress):
+        self.soa = soa
+        self.__scheme = soa.scheme
+        self.__account_name = soa.account
+        self.__container_name = soa.container
+        self.__prefix_name = soa.prefix
+
+        # Reconstruct base URI in normalized format
+        self.__base_uri = self.soa.prefix_free_uri + self.__prefix_name
+        _logger.info(
+            f'Initialized AzureBlobStore with base URI: {self.__base_uri},',
+            f'account: {self.__account_name}, container: {self.__container_name}, prefix: {self.__prefix_name}',
+        )
+
+    def client(self) -> 'BlobServiceClient':
+        """Return the Azure Blob Storage client."""
+        client_dict: dict[str, 'BlobServiceClient'] = env.Env.get().get_client('azure_blob')
+        with client_lock:
+            uri = self.soa.container_free_uri
+            if uri not in client_dict:
+                storage_account_name = Config.get().get_string_value('storage_account_name', section='azure')
+                storage_account_key = Config.get().get_string_value('storage_account_key', section='azure')
+                if (storage_account_name is None) != (storage_account_key is None):
+                    raise excs.Error(
+                        "Azure 'storage_account_name' and 'storage_account_key' must be specified together."
+                    )
+                if storage_account_name is None or storage_account_name != self.__account_name:
+                    # Attempt a connection to a public resource, with no account key
+                    client_dict[uri] = self.create_client(endpoint_url=uri)
+                else:
+                    client_dict[uri] = self.create_client(
+                        endpoint_url=uri, account_name=self.__account_name, account_key=storage_account_key
+                    )
+            return client_dict[uri]
+
+    @property
+    def account_name(self) -> str:
+        """Return the storage account name."""
+        return self.__account_name
+
+    @property
+    def container_name(self) -> str:
+        """Return the container name from the base URI."""
+        return self.__container_name
+
+    @property
+    def prefix(self) -> str:
+        """Return the prefix from the base URI."""
+        return self.__prefix_name
+
+    def validate(self, error_col_name: str) -> str | None:
+        """
+        Checks if the URI exists and is accessible.
+
+        Returns:
+            str: The base URI if the container exists and is accessible, None otherwise.
+        """
+        try:
+            container_client = self.client().get_container_client(self.container_name)
+            # Check if container exists by trying to get its properties
+            container_client.get_container_properties()
+            return self.__base_uri
+        except AzureError as e:
+            self.handle_azure_error(e, self.container_name, f'validate container {error_col_name}')
+        return None
+
+    def copy_object_to_local_file(self, src_path: str, dest_path: Path) -> None:
+        """Copies a blob to a local file. Thread safe."""
+        try:
+            blob_client = self.client().get_blob_client(container=self.container_name, blob=self.prefix + src_path)
+            with open(dest_path, 'wb') as download_file:
+                download_stream = blob_client.download_blob()
+                download_file.write(download_stream.readall())
+        except AzureError as e:
+            self.handle_azure_error(e, self.container_name, f'download file {src_path}')
+            raise
+
+    # TODO: utils package should not include back-references to `Column`
+    def copy_local_file(self, col: 'Column', src_path: Path) -> str:
+        """Copy a local file to Azure Blob Storage, and return its new URL"""
+        prefix, filename = ObjectPath.create_prefix_raw(
+            col.get_tbl().id, col.id, col.get_tbl().version, ext=src_path.suffix
+        )
+        blob_name = f'{self.prefix}{prefix}/{filename}'
+        new_file_uri = f'{self.__base_uri}{prefix}/{filename}'
+
+        try:
+            blob_client = self.client().get_blob_client(container=self.container_name, blob=blob_name)
+            with open(src_path, 'rb') as data:
+                blob_client.upload_blob(data, overwrite=True)
+            _logger.debug(f'Media Storage: copied {src_path} to {new_file_uri}')
+            return new_file_uri
+        except AzureError as e:
+            self.handle_azure_error(e, self.container_name, f'upload file {src_path}')
+            raise
+
+    def _get_filtered_blobs(
+        self, tbl_id: uuid.UUID | None, tbl_version: int | None = None
+    ) -> Iterator['BlobProperties']:
+        """Private method to get filtered blobs for a table, optionally filtered by version.
+
+        Args:
+            tbl_id: Table UUID to filter by
+            tbl_version: Optional table version to filter by
+
+        Returns:
+            Iterator over blob objects matching the criteria
+        """
+        # Use ObjectPath to construct the prefix for this table
+        if tbl_id is None:
+            prefix = self.prefix
+            assert tbl_version is None, 'tbl_version must be None if tbl_id is None'
+        else:
+            table_prefix = ObjectPath.table_prefix(tbl_id)
+            prefix = f'{self.prefix}{table_prefix}/'
+
+        try:
+            container_client = self.client().get_container_client(self.container_name)
+
+            blob_iterator: Iterator['BlobProperties']
+            if tbl_version is None:
+                # Return all blobs with the table prefix
+                blob_iterator = container_client.list_blobs(name_starts_with=prefix)
+            else:
+                # Filter by both table_id and table_version using the ObjectPath pattern
+                # Pattern: tbl_id_col_id_version_uuid
+                version_pattern = re.compile(
+                    rf'{re.escape(table_prefix)}_\d+_{re.escape(str(tbl_version))}_[0-9a-fA-F]+.*'
+                )
+                # Get all blobs with the prefix and filter by version pattern
+                all_blobs = container_client.list_blobs(name_starts_with=prefix)
+                blob_iterator = (blob for blob in all_blobs if version_pattern.match(blob.name.split('/')[-1]))
+
+            return blob_iterator
+
+        except AzureError as e:
+            self.handle_azure_error(e, self.container_name, f'setup iterator {self.prefix}')
+            raise
+
+    def count(self, tbl_id: uuid.UUID | None, tbl_version: int | None = None) -> int:
+        """Count the number of files belonging to tbl_id. If tbl_version is not None,
+        count only those files belonging to the specified tbl_version.
+
+        Args:
+            tbl_id: Table UUID to count blobs for
+            tbl_version: Optional table version to filter by
+
+        Returns:
+            Number of blobs matching the criteria
+        """
+        blob_iterator = self._get_filtered_blobs(tbl_id, tbl_version)
+        return sum(1 for _ in blob_iterator)
+
+    def delete(self, tbl_id: uuid.UUID, tbl_version: int | None = None) -> int:
+        """Delete all files belonging to tbl_id. If tbl_version is not None, delete
+        only those files belonging to the specified tbl_version.
+
+        Args:
+            tbl_id: Table UUID to delete blobs for
+            tbl_version: Optional table version to filter by
+
+        Returns:
+            Number of blobs deleted
+        """
+        assert tbl_id is not None
+        blob_iterator = self._get_filtered_blobs(tbl_id, tbl_version)
+        total_deleted = 0
+
+        try:
+            container_client = self.client().get_container_client(self.container_name)
+
+            for blob in blob_iterator:
+                # TODO: Figure out now to properly use batch method delete_blobs(), it doesn't seem to work properly
+                container_client.delete_blob(blob.name)
+                total_deleted += 1
+
+            # print(f"Deleted {total_deleted} blobs from container '{self.container_name}'.")
+            return total_deleted
+
+        except AzureError as e:
+            self.handle_azure_error(e, self.container_name, f'deleting with {self.prefix}')
+            raise
+
+    def list_objects(self, return_uri: bool, n_max: int = 10) -> list[str]:
+        """Return a list of objects found in the specified destination bucket.
+        Each returned object includes the full set of prefixes.
+        if return_uri is True, full URI's are returned; otherwise, just the object keys.
+        """
+        p = self.soa.prefix_free_uri if return_uri else ''
+        r: list[str] = []
+        try:
+            blob_iterator = self._get_filtered_blobs(tbl_id=None, tbl_version=None)
+            for blob in blob_iterator:
+                r.append(f'{p}{blob.name}')
+                if len(r) >= n_max:
+                    return r
+
+        except AzureError as e:
+            self.handle_azure_error(e, self.__container_name, f'list objects from {self.__base_uri}')
+        return r
+
+    @classmethod
+    def handle_azure_error(
+        cls, e: 'AzureError', container_name: str, operation: str = '', *, ignore_404: bool = False
+    ) -> None:
+        from azure.core.exceptions import ClientAuthenticationError, HttpResponseError, ResourceNotFoundError
+
+        if ignore_404 and isinstance(e, ResourceNotFoundError):
+            return
+
+        if isinstance(e, ResourceNotFoundError):
+            raise excs.Error(f'Container {container_name} or blob not found during {operation}: {str(e)!r}')
+        elif isinstance(e, ClientAuthenticationError):
+            raise excs.Error(f'Authentication failed for container {container_name} during {operation}: {str(e)!r}')
+        elif isinstance(e, HttpResponseError):
+            if e.status_code == 403:
+                raise excs.Error(f'Access denied to container {container_name} during {operation}: {str(e)!r}')
+            elif e.status_code == 412:
+                raise excs.Error(f'Precondition failed for container {container_name} during {operation}: {str(e)!r}')
+            else:
+                raise excs.Error(
+                    f'HTTP error during {operation} in container {container_name}: {e.status_code} - {str(e)!r}'
+                )
+        else:
+            raise excs.Error(f'Error during {operation} in container {container_name}: {str(e)!r}')
+
+    @classmethod
+    def create_client(
+        cls, endpoint_url: str, account_name: str | None = None, account_key: str | None = None
+    ) -> 'BlobServiceClient':
+        from azure.core.credentials import AzureNamedKeyCredential
+        from azure.storage.blob import BlobServiceClient  # TODO: Use azure.storage.blob.aio instead
+
+        assert (account_name is None) == (account_key is None)
+        try:
+            #  e.g. endpoint_url: str = f'https://{account_name}.blob.core.windows.net'
+            assert endpoint_url is not None, 'No Azure Storage account information provided'
+
+            # Use empty SAS token for anonymous authentication
+            credential = None
+            if account_name is not None:
+                credential = AzureNamedKeyCredential(name=account_name, key=account_key)
+            return BlobServiceClient(
+                account_url=endpoint_url,
+                credential=credential,
+                max_single_get_size=(32 * 2**20),
+                max_chunk_get_size=(4 * 2**20),
+                connection_timeout=15,
+                read_timeout=30,
+            )
+        except Exception as e:
+            raise excs.Error(f'Failed to create Azure Blob Storage client: {str(e)!r}') from e

pixeltable/utils/code.py

@@ -1,12 +1,11 @@
 import types
-from typing import Optional
 
 from pixeltable.func import Function
 
 # Utilities related to the organization of the Pixeltable codebase.
 
 
-def local_public_names(mod_name: str, exclude: Optional[list[str]] = None) -> list[str]:
+def local_public_names(mod_name: str, exclude: list[str] | None = None) -> list[str]:
     """
     Returns a list of all functions and submodules that are local to the specified module and are
     publicly accessible. Intended to facilitate implementation of module __dir__() methods for
@@ -21,7 +20,8 @@ def local_public_names(mod_name: str, exclude: Optional[list[str]] = None) -> li
     for obj in mod.__dict__.values():
         if isinstance(obj, Function):
             # Pixeltable function
-            names.append(obj.name)
+            if not obj.name.startswith('_'):
+                names.append(obj.name)
         elif isinstance(obj, types.FunctionType):
             # Python function
             if obj.__module__ == mod.__name__ and not obj.__name__.startswith('_'):

pixeltable/utils/console_output.py

@@ -1,6 +1,8 @@
 import logging
 from typing import TextIO
 
+from pixeltable import exceptions as excs
+
 
 def map_level(verbosity: int) -> int:
     """
@@ -19,7 +21,8 @@ def map_level(verbosity: int) -> int:
         return logging.INFO
     if verbosity == 2:
         return logging.DEBUG
-    return logging.INFO
+
+    raise excs.Error(f'Invalid verbosity level: {verbosity}')
 
 
 class ConsoleOutputHandler(logging.StreamHandler):

pixeltable/utils/coroutine.py

@@ -1,10 +1,10 @@
 import asyncio
 import threading
-from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Coroutine, TypeVar
 
-T = TypeVar('T')
+from pixeltable.env import Env
 
+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., as an embedding function as part of a similarity lookup). Once the scheduler is fully
@@ -15,27 +15,10 @@ 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() -> T:
-        new_loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(new_loop)
-        try:
-            return new_loop.run_until_complete(coroutine)
-        finally:
-            new_loop.close()
-
-    try:
-        loop = asyncio.get_running_loop()
-    except RuntimeError:
-        # No event loop; just call `asyncio.run()`
-        return asyncio.run(coroutine)
+    loop = Env.get().event_loop
 
     if threading.current_thread() is threading.main_thread():
-        if not loop.is_running():
-            return loop.run_until_complete(coroutine)
-        else:
-            with ThreadPoolExecutor() as pool:
-                future = pool.submit(run_in_new_loop)
-                return future.result(timeout=timeout)
+        return loop.run_until_complete(coroutine)
     else:
-        return asyncio.run_coroutine_threadsafe(coroutine, loop).result()
+        # Not in main thread, use run_coroutine_threadsafe
+        return asyncio.run_coroutine_threadsafe(coroutine, loop).result(timeout)