| """ |
| Base classes and types for data sources. |
| |
| This module defines the abstract base class for all data sources and |
| common types used throughout the data sources subsystem. |
| """ |
|
|
| from abc import ABC, abstractmethod |
| from dataclasses import dataclass, field |
| from enum import Enum |
| from typing import Any, Dict, Iterator, List, Optional |
| import logging |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class SourceType(Enum): |
| """Enumeration of supported data source types.""" |
| FILE = "file" |
| URL = "url" |
| GOOGLE_DRIVE = "google_drive" |
| DROPBOX = "dropbox" |
| S3 = "s3" |
| HUGGINGFACE = "huggingface" |
| GOOGLE_SHEETS = "google_sheets" |
| DATABASE = "database" |
|
|
|
|
| @dataclass |
| class SourceConfig: |
| """ |
| Configuration for a data source. |
| |
| This dataclass holds the parsed configuration for a single data source, |
| including type-specific settings and common options. |
| |
| Attributes: |
| source_type: The type of data source |
| source_id: Unique identifier for this source (auto-generated if not provided) |
| config: The raw configuration dictionary for this source |
| enabled: Whether this source is enabled |
| """ |
| source_type: SourceType |
| source_id: str |
| config: Dict[str, Any] |
| enabled: bool = True |
|
|
| @classmethod |
| def from_dict(cls, config_dict: Dict[str, Any], index: int = 0) -> "SourceConfig": |
| """ |
| Create a SourceConfig from a configuration dictionary. |
| |
| Args: |
| config_dict: Dictionary containing source configuration |
| index: Index in the sources list (used for auto-generated ID) |
| |
| Returns: |
| SourceConfig instance |
| |
| Raises: |
| ValueError: If the source type is invalid or required fields are missing |
| """ |
| type_str = config_dict.get("type") |
| if not type_str: |
| raise ValueError("Data source configuration must include 'type' field") |
|
|
| try: |
| source_type = SourceType(type_str) |
| except ValueError: |
| valid_types = [t.value for t in SourceType] |
| raise ValueError( |
| f"Invalid data source type '{type_str}'. " |
| f"Valid types are: {', '.join(valid_types)}" |
| ) |
|
|
| |
| source_id = config_dict.get("id") or config_dict.get("source_id") |
| if not source_id: |
| |
| if source_type == SourceType.FILE: |
| path = config_dict.get("path", "") |
| source_id = f"file_{index}_{path.replace('/', '_').replace('.', '_')}" |
| elif source_type == SourceType.URL: |
| url = config_dict.get("url", "") |
| |
| source_id = f"url_{index}_{url.split('/')[-1][:30]}" |
| else: |
| source_id = f"{source_type.value}_{index}" |
|
|
| enabled = config_dict.get("enabled", True) |
|
|
| return cls( |
| source_type=source_type, |
| source_id=source_id, |
| config=config_dict, |
| enabled=enabled |
| ) |
|
|
|
|
| class DataSource(ABC): |
| """ |
| Abstract base class for all data sources. |
| |
| Each data source implementation must provide methods for: |
| - Identifying the source |
| - Checking availability |
| - Reading items (with optional partial reading support) |
| - Reporting total item count |
| |
| Thread Safety: |
| Implementations should be thread-safe for concurrent read operations. |
| Write operations (if any) should be protected by appropriate locking. |
| """ |
|
|
| def __init__(self, config: SourceConfig): |
| """ |
| Initialize the data source. |
| |
| Args: |
| config: Source configuration |
| """ |
| self._config = config |
| self._source_id = config.source_id |
| self._raw_config = config.config |
|
|
| @property |
| def source_id(self) -> str: |
| """Get the unique identifier for this source.""" |
| return self._source_id |
|
|
| @property |
| def source_type(self) -> SourceType: |
| """Get the type of this source.""" |
| return self._config.source_type |
|
|
| @property |
| def config(self) -> Dict[str, Any]: |
| """Get the raw configuration dictionary.""" |
| return self._raw_config |
|
|
| @abstractmethod |
| def get_source_id(self) -> str: |
| """ |
| Get the unique identifier for this data source. |
| |
| Returns: |
| String identifier unique within the DataSourceManager |
| """ |
| pass |
|
|
| @abstractmethod |
| def is_available(self) -> bool: |
| """ |
| Check if the data source is available and accessible. |
| |
| This method should verify that: |
| - Required dependencies are installed |
| - Credentials are valid (if applicable) |
| - The source location exists and is readable |
| |
| Returns: |
| True if the source is ready to read from |
| """ |
| pass |
|
|
| @abstractmethod |
| def read_items( |
| self, |
| start: int = 0, |
| count: Optional[int] = None |
| ) -> Iterator[Dict[str, Any]]: |
| """ |
| Read items from the data source. |
| |
| This method yields dictionaries representing annotation items. |
| For sources that support partial reading, the start and count |
| parameters allow fetching specific ranges of items. |
| |
| Args: |
| start: Index of the first item to read (0-based) |
| count: Maximum number of items to read (None = all remaining) |
| |
| Yields: |
| Dictionary containing item data with at least id_key field |
| |
| Raises: |
| RuntimeError: If the source is not available |
| IOError: If reading fails |
| """ |
| pass |
|
|
| @abstractmethod |
| def get_total_count(self) -> Optional[int]: |
| """ |
| Get the total number of items in the source. |
| |
| Returns: |
| Total item count, or None if unknown (e.g., streaming source) |
| """ |
| pass |
|
|
| @abstractmethod |
| def supports_partial_reading(self) -> bool: |
| """ |
| Check if this source supports reading partial ranges. |
| |
| Partial reading allows loading data in chunks, which is useful |
| for large datasets or incremental annotation workflows. |
| |
| Returns: |
| True if read_items() supports start/count parameters |
| """ |
| pass |
|
|
| def validate_config(self) -> List[str]: |
| """ |
| Validate the source configuration. |
| |
| Override this method to add source-specific validation. |
| |
| Returns: |
| List of validation error messages (empty if valid) |
| """ |
| return [] |
|
|
| def get_status(self) -> Dict[str, Any]: |
| """ |
| Get the current status of this data source. |
| |
| Returns: |
| Dictionary containing status information: |
| - source_id: The source identifier |
| - source_type: The source type |
| - available: Whether the source is available |
| - total_count: Total item count (or None) |
| - supports_partial: Whether partial reading is supported |
| """ |
| return { |
| "source_id": self.source_id, |
| "source_type": self.source_type.value, |
| "available": self.is_available(), |
| "total_count": self.get_total_count(), |
| "supports_partial": self.supports_partial_reading(), |
| } |
|
|
| def refresh(self) -> bool: |
| """ |
| Refresh the data source (re-fetch from remote, re-validate, etc.). |
| |
| Override this method for sources that cache data or need |
| periodic refresh. The default implementation does nothing. |
| |
| Returns: |
| True if refresh was successful |
| """ |
| return True |
|
|
| def close(self) -> None: |
| """ |
| Close the data source and release any resources. |
| |
| Override this method for sources that hold resources like |
| database connections or file handles. |
| """ |
| pass |
|
|
| def __enter__(self) -> "DataSource": |
| """Context manager entry.""" |
| return self |
|
|
| def __exit__(self, exc_type, exc_val, exc_tb) -> None: |
| """Context manager exit.""" |
| self.close() |
|
|
| def __repr__(self) -> str: |
| return f"{self.__class__.__name__}(source_id={self.source_id!r})" |
|
|
