datachain 0.34.6__py3-none-any.whl → 0.35.0__py3-none-any.whl

datachain/data_storage/sqlite.py

@@ -1,18 +1,11 @@
 import logging
 import os
 import sqlite3
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from contextlib import contextmanager
 from functools import cached_property, wraps
 from time import sleep
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    ClassVar,
-    Optional,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, ClassVar, Union
 
 import sqlalchemy
 from sqlalchemy import (
@@ -105,8 +98,8 @@ def retry_sqlite_locks(func):
 
 
 def get_db_file_in_memory(
-    db_file: Optional[str] = None, in_memory: bool = False
-) -> Optional[str]:
+    db_file: str | None = None, in_memory: bool = False
+) -> str | None:
     """Get in-memory db_file and check that conflicting arguments are not provided."""
     if in_memory:
         if db_file and db_file != ":memory:":
@@ -119,7 +112,7 @@ class SQLiteDatabaseEngine(DatabaseEngine):
     dialect = sqlite_dialect
 
     db: sqlite3.Connection
-    db_file: Optional[str]
+    db_file: str | None
     is_closed: bool
 
     def __init__(
@@ -127,8 +120,8 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         engine: "Engine",
         metadata: "MetaData",
         db: sqlite3.Connection,
-        db_file: Optional[str] = None,
-        max_variable_number: Optional[int] = 999,
+        db_file: str | None = None,
+        max_variable_number: int | None = 999,
     ):
         self.engine = engine
         self.metadata = metadata
@@ -138,12 +131,12 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         self.max_variable_number = max_variable_number
 
     @classmethod
-    def from_db_file(cls, db_file: Optional[str] = None) -> "SQLiteDatabaseEngine":
+    def from_db_file(cls, db_file: str | None = None) -> "SQLiteDatabaseEngine":
         return cls(*cls._connect(db_file=db_file))
 
     @staticmethod
     def _connect(
-        db_file: Optional[str] = None,
+        db_file: str | None = None,
     ) -> tuple["Engine", "MetaData", sqlite3.Connection, str, int]:
         try:
             if db_file == ":memory:":
@@ -232,7 +225,7 @@ class SQLiteDatabaseEngine(DatabaseEngine):
     def execute(
         self,
         query,
-        cursor: Optional[sqlite3.Cursor] = None,
+        cursor: sqlite3.Cursor | None = None,
         conn=None,
     ) -> sqlite3.Cursor:
         if self.is_closed:
@@ -251,7 +244,7 @@ class SQLiteDatabaseEngine(DatabaseEngine):
 
     @retry_sqlite_locks
     def executemany(
-        self, query, params, cursor: Optional[sqlite3.Cursor] = None, conn=None
+        self, query, params, cursor: sqlite3.Cursor | None = None, conn=None
     ) -> sqlite3.Cursor:
         if cursor:
             return cursor.executemany(self.compile(query).string, params)
@@ -351,13 +344,13 @@ class SQLiteMetastore(AbstractDBMetastore):
 
     META_TABLE = "meta"
 
-    db: "SQLiteDatabaseEngine"
+    db: SQLiteDatabaseEngine
 
     def __init__(
         self,
-        uri: Optional[StorageURI] = None,
-        db: Optional["SQLiteDatabaseEngine"] = None,
-        db_file: Optional[str] = None,
+        uri: StorageURI | None = None,
+        db: SQLiteDatabaseEngine | None = None,
+        db_file: str | None = None,
         in_memory: bool = False,
     ):
         uri = uri or StorageURI("")
@@ -384,7 +377,7 @@ class SQLiteMetastore(AbstractDBMetastore):
 
     def clone(
         self,
-        uri: Optional[StorageURI] = None,
+        uri: StorageURI | None = None,
         use_new_connection: bool = False,
     ) -> "SQLiteMetastore":
         uri = uri or StorageURI("")
@@ -582,15 +575,15 @@ class SQLiteWarehouse(AbstractWarehouse):
     This is currently used for the local cli.
     """
 
-    db: "SQLiteDatabaseEngine"
+    db: SQLiteDatabaseEngine
 
     # Cache for our defined column types to dialect specific TypeEngine relations
     _col_python_type: ClassVar[dict[type, "TypeEngine"]] = {}
 
     def __init__(
         self,
-        db: Optional["SQLiteDatabaseEngine"] = None,
-        db_file: Optional[str] = None,
+        db: SQLiteDatabaseEngine | None = None,
+        db_file: str | None = None,
         in_memory: bool = False,
     ):
         self.schema: DefaultSchema = DefaultSchema()
@@ -645,7 +638,7 @@ class SQLiteWarehouse(AbstractWarehouse):
             only=filter_tables,
         )
 
-    def is_ready(self, timeout: Optional[int] = None) -> bool:
+    def is_ready(self, timeout: int | None = None) -> bool:
         return True
 
     def create_dataset_rows_table(
@@ -791,7 +784,7 @@ class SQLiteWarehouse(AbstractWarehouse):
         self,
         table: Table,
         query: Select,
-        progress_cb: Optional[Callable[[int], None]] = None,
+        progress_cb: Callable[[int], None] | None = None,
     ) -> None:
         col_id = (
             query.selected_columns.sys__id

datachain/data_storage/warehouse.py

@@ -4,8 +4,8 @@ import posixpath
 import random
 import string
 from abc import ABC, abstractmethod
-from collections.abc import Generator, Iterable, Iterator, Sequence
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+from collections.abc import Callable, Generator, Iterable, Iterator, Sequence
+from typing import TYPE_CHECKING, Any, Union
 from urllib.parse import urlparse
 
 import attrs
@@ -174,12 +174,12 @@ class AbstractWarehouse(ABC, Serializable):
     #
 
     @abstractmethod
-    def is_ready(self, timeout: Optional[int] = None) -> bool: ...
+    def is_ready(self, timeout: int | None = None) -> bool: ...
 
     def dataset_rows(
         self,
         dataset: DatasetRecord,
-        version: Optional[str] = None,
+        version: str | None = None,
         column: str = "file",
     ):
         version = version or dataset.latest_version
@@ -424,7 +424,7 @@ class AbstractWarehouse(ABC, Serializable):
 
     def dataset_stats(
         self, dataset: DatasetRecord, version: str
-    ) -> tuple[Optional[int], Optional[int]]:
+    ) -> tuple[int | None, int | None]:
         """
         Returns tuple with dataset stats: total number of rows and total dataset size.
         """
@@ -549,7 +549,7 @@ class AbstractWarehouse(ABC, Serializable):
         dr = dataset_rows
         columns = [c.name for c in query.selected_columns]
         for row in self.db.execute(query):
-            d = dict(zip(columns, row))
+            d = dict(zip(columns, row, strict=False))
             yield Node(**{dr.without_object(k): v for k, v in d.items()})
 
     def get_dirs_by_parent_path(
@@ -786,7 +786,7 @@ class AbstractWarehouse(ABC, Serializable):
     def size(
         self,
         dataset_rows: "DataTable",
-        node: Union[Node, dict[str, Any]],
+        node: Node | dict[str, Any],
         count_files: bool = False,
     ) -> tuple[int, int]:
         """
@@ -828,10 +828,10 @@ class AbstractWarehouse(ABC, Serializable):
         self,
         dataset_rows: "DataTable",
         parent_path: str,
-        fields: Optional[Sequence[str]] = None,
-        type: Optional[str] = None,
+        fields: Sequence[str] | None = None,
+        type: str | None = None,
         conds=None,
-        order_by: Optional[Union[str, list[str]]] = None,
+        order_by: str | list[str] | None = None,
         include_subobjects: bool = True,
     ) -> sa.Select:
         if not conds:
@@ -869,7 +869,7 @@ class AbstractWarehouse(ABC, Serializable):
         self,
         dataset_rows: "DataTable",
         node: Node,
-        sort: Union[list[str], str, None] = None,
+        sort: list[str] | str | None = None,
         include_subobjects: bool = True,
     ) -> Iterator[NodeWithPath]:
         """
@@ -927,7 +927,7 @@ class AbstractWarehouse(ABC, Serializable):
     def create_udf_table(
         self,
         columns: Sequence["sa.Column"] = (),
-        name: Optional[str] = None,
+        name: str | None = None,
     ) -> sa.Table:
         """
         Create a temporary table for storing custom signals generated by a UDF.
@@ -948,7 +948,7 @@ class AbstractWarehouse(ABC, Serializable):
         self,
         table: sa.Table,
         query: sa.Select,
-        progress_cb: Optional[Callable[[int], None]] = None,
+        progress_cb: Callable[[int], None] | None = None,
     ) -> None:
         """
         Copy the results of a query into a table.

datachain/dataset.py

@@ -3,13 +3,7 @@ import json
 from dataclasses import dataclass, fields
 from datetime import datetime
 from functools import cached_property
-from typing import (
-    Any,
-    NewType,
-    Optional,
-    TypeVar,
-    Union,
-)
+from typing import Any, NewType, TypeVar
 from urllib.parse import urlparse
 
 from packaging.specifiers import SpecifierSet
@@ -43,7 +37,7 @@ DATASET_NAME_REPLACEMENT_CHAR = "_"
 StorageURI = NewType("StorageURI", str)
 
 
-def parse_dataset_uri(uri: str) -> tuple[str, Optional[str]]:
+def parse_dataset_uri(uri: str) -> tuple[str, str | None]:
     """
     Parse dataser uri to extract name and version out of it (if version is defined)
     Example:
@@ -65,7 +59,7 @@ def parse_dataset_uri(uri: str) -> tuple[str, Optional[str]]:
 
 
 def create_dataset_uri(
-    name: str, namespace: str, project: str, version: Optional[str] = None
+    name: str, namespace: str, project: str, version: str | None = None
 ) -> str:
     """
     Creates a dataset uri based on namespace, project, dataset name and optionally
@@ -81,7 +75,7 @@ def create_dataset_uri(
     return uri
 
 
-def parse_dataset_name(name: str) -> tuple[Optional[str], Optional[str], str]:
+def parse_dataset_name(name: str) -> tuple[str | None, str | None, str]:
     """Parses dataset name and returns namespace, project and name"""
     if not name:
         raise InvalidDatasetNameError("Name must be defined to parse it")
@@ -111,7 +105,7 @@ class DatasetDependency:
     name: str
     version: str
     created_at: datetime
-    dependencies: list[Optional["DatasetDependency"]]
+    dependencies: list["DatasetDependency | None"]
 
     @property
     def dataset_name(self) -> str:
@@ -131,12 +125,12 @@ class DatasetDependency:
         namespace_name: str,
         project_name: str,
         id: int,
-        dataset_id: Optional[int],
-        dataset_version_id: Optional[int],
-        dataset_name: Optional[str],
-        dataset_version: Optional[str],
-        dataset_version_created_at: Optional[datetime],
-    ) -> Optional["DatasetDependency"]:
+        dataset_id: int | None,
+        dataset_version_id: int | None,
+        dataset_name: str | None,
+        dataset_version: str | None,
+        dataset_version_created_at: datetime | None,
+    ) -> "DatasetDependency | None":
         from datachain.lib.listing import is_listing_dataset
 
         if not dataset_id:
@@ -198,17 +192,17 @@ class DatasetVersion:
     status: int
     feature_schema: dict
     created_at: datetime
-    finished_at: Optional[datetime]
+    finished_at: datetime | None
     error_message: str
     error_stack: str
     script_output: str
-    schema: dict[str, Union[SQLType, type[SQLType]]]
-    num_objects: Optional[int]
-    size: Optional[int]
-    _preview_data: Optional[Union[str, list[dict]]]
+    schema: dict[str, SQLType | type[SQLType]]
+    num_objects: int | None
+    size: int | None
+    _preview_data: str | list[dict] | None
     sources: str = ""
     query_script: str = ""
-    job_id: Optional[str] = None
+    job_id: str | None = None
 
     @classmethod
     def parse(  # noqa: PLR0913
@@ -218,19 +212,19 @@ class DatasetVersion:
         dataset_id: int,
         version: str,
         status: int,
-        feature_schema: Optional[str],
+        feature_schema: str | None,
         created_at: datetime,
-        finished_at: Optional[datetime],
+        finished_at: datetime | None,
         error_message: str,
         error_stack: str,
         script_output: str,
-        num_objects: Optional[int],
-        size: Optional[int],
-        preview: Optional[Union[str, list[dict]]],
-        schema: dict[str, Union[SQLType, type[SQLType]]],
+        num_objects: int | None,
+        size: int | None,
+        preview: str | list[dict] | None,
+        schema: dict[str, SQLType | type[SQLType]],
         sources: str = "",
         query_script: str = "",
-        job_id: Optional[str] = None,
+        job_id: str | None = None,
     ):
         return cls(
             id,
@@ -292,7 +286,7 @@ class DatasetVersion:
         }
 
     @cached_property
-    def preview(self) -> Optional[list[dict]]:
+    def preview(self) -> list[dict] | None:
         if isinstance(self._preview_data, str):
             return json.loads(self._preview_data)
         return self._preview_data if self._preview_data else None
@@ -313,13 +307,13 @@ class DatasetListVersion:
     version: str
     status: int
     created_at: datetime
-    finished_at: Optional[datetime]
+    finished_at: datetime | None
     error_message: str
     error_stack: str
-    num_objects: Optional[int]
-    size: Optional[int]
+    num_objects: int | None
+    size: int | None
     query_script: str = ""
-    job_id: Optional[str] = None
+    job_id: str | None = None
 
     @classmethod
     def parse(
@@ -330,13 +324,13 @@ class DatasetListVersion:
         version: str,
         status: int,
         created_at: datetime,
-        finished_at: Optional[datetime],
+        finished_at: datetime | None,
         error_message: str,
         error_stack: str,
-        num_objects: Optional[int],
-        size: Optional[int],
+        num_objects: int | None,
+        size: int | None,
         query_script: str = "",
-        job_id: Optional[str] = None,
+        job_id: str | None = None,
         **kwargs,
     ):
         return cls(
@@ -368,14 +362,14 @@ class DatasetRecord:
     id: int
     name: str
     project: Project
-    description: Optional[str]
+    description: str | None
     attrs: list[str]
-    schema: dict[str, Union[SQLType, type[SQLType]]]
+    schema: dict[str, SQLType | type[SQLType]]
     feature_schema: dict
     versions: list[DatasetVersion]
     status: int = DatasetStatus.CREATED
-    created_at: Optional[datetime] = None
-    finished_at: Optional[datetime] = None
+    created_at: datetime | None = None
+    finished_at: datetime | None = None
     error_message: str = ""
     error_stack: str = ""
     script_output: str = ""
@@ -388,7 +382,7 @@ class DatasetRecord:
     @staticmethod
     def parse_schema(
         ct: dict[str, Any],
-    ) -> dict[str, Union[SQLType, type[SQLType]]]:
+    ) -> dict[str, SQLType | type[SQLType]]:
         return {
             c_name: NAME_TYPES_MAPPING[c_type["type"]].from_dict(c_type)  # type: ignore [attr-defined]
             for c_name, c_type in ct.items()
@@ -409,23 +403,23 @@ class DatasetRecord:
         namespace_id: int,
         namespace_uuid: str,
         namespace_name: str,
-        namespace_description: Optional[str],
+        namespace_description: str | None,
         namespace_created_at: datetime,
         project_id: int,
         project_uuid: str,
         project_name: str,
-        project_description: Optional[str],
+        project_description: str | None,
         project_created_at: datetime,
         project_namespace_id: int,
         dataset_id: int,
         dataset_project_id: int,
         name: str,
-        description: Optional[str],
+        description: str | None,
         attrs: str,
         status: int,
-        feature_schema: Optional[str],
+        feature_schema: str | None,
         created_at: datetime,
-        finished_at: Optional[datetime],
+        finished_at: datetime | None,
         error_message: str,
         error_stack: str,
         script_output: str,
@@ -437,19 +431,19 @@ class DatasetRecord:
         version_dataset_id: int,
         version: str,
         version_status: int,
-        version_feature_schema: Optional[str],
+        version_feature_schema: str | None,
         version_created_at: datetime,
-        version_finished_at: Optional[datetime],
+        version_finished_at: datetime | None,
         version_error_message: str,
         version_error_stack: str,
         version_script_output: str,
-        version_num_objects: Optional[int],
-        version_size: Optional[int],
-        version_preview: Optional[str],
-        version_sources: Optional[str],
-        version_query_script: Optional[str],
+        version_num_objects: int | None,
+        version_size: int | None,
+        version_preview: str | None,
+        version_sources: str | None,
+        version_query_script: str | None,
         version_schema: str,
-        version_job_id: Optional[str] = None,
+        version_job_id: str | None = None,
     ) -> "DatasetRecord":
         attrs_lst: list[str] = json.loads(attrs) if attrs else []
         schema_dct: dict[str, Any] = json.loads(schema) if schema else {}
@@ -527,7 +521,7 @@ class DatasetRecord:
     def full_name(self) -> str:
         return f"{self.project.namespace.name}.{self.project.name}.{self.name}"
 
-    def get_schema(self, version: str) -> dict[str, Union[SQLType, type[SQLType]]]:
+    def get_schema(self, version: str) -> dict[str, SQLType | type[SQLType]]:
         return self.get_version(version).schema if version else self.schema
 
     def update(self, **kwargs):
@@ -649,7 +643,7 @@ class DatasetRecord:
         """Returns latest version of a dataset"""
         return max(self.versions).version
 
-    def latest_major_version(self, major: int) -> Optional[str]:
+    def latest_major_version(self, major: int) -> str | None:
         """
         Returns latest specific major version, e.g if dataset has versions:
             - 1.4.1
@@ -664,7 +658,7 @@ class DatasetRecord:
             return None
         return max(versions).version
 
-    def latest_compatible_version(self, version_spec: str) -> Optional[str]:
+    def latest_compatible_version(self, version_spec: str) -> str | None:
         """
         Returns the latest version that matches the given version specifier.
 
@@ -711,10 +705,10 @@ class DatasetListRecord:
     id: int
     name: str
     project: Project
-    description: Optional[str]
+    description: str | None
     attrs: list[str]
     versions: list[DatasetListVersion]
-    created_at: Optional[datetime] = None
+    created_at: datetime | None = None
 
     @classmethod
     def parse(  # noqa: PLR0913
@@ -722,17 +716,17 @@ class DatasetListRecord:
         namespace_id: int,
         namespace_uuid: str,
         namespace_name: str,
-        namespace_description: Optional[str],
+        namespace_description: str | None,
         namespace_created_at: datetime,
         project_id: int,
         project_uuid: str,
         project_name: str,
-        project_description: Optional[str],
+        project_description: str | None,
         project_created_at: datetime,
         project_namespace_id: int,
         dataset_id: int,
         name: str,
-        description: Optional[str],
+        description: str | None,
         attrs: str,
         created_at: datetime,
         version_id: int,
@@ -741,13 +735,13 @@ class DatasetListRecord:
         version: str,
         version_status: int,
         version_created_at: datetime,
-        version_finished_at: Optional[datetime],
+        version_finished_at: datetime | None,
         version_error_message: str,
         version_error_stack: str,
-        version_num_objects: Optional[int],
-        version_size: Optional[int],
-        version_query_script: Optional[str],
-        version_job_id: Optional[str] = None,
+        version_num_objects: int | None,
+        version_size: int | None,
+        version_query_script: str | None,
+        version_job_id: str | None = None,
     ) -> "DatasetListRecord":
         attrs_lst: list[str] = json.loads(attrs) if attrs else []
 

datachain/delta.py

@@ -1,7 +1,7 @@
 from collections.abc import Sequence
 from copy import copy
 from functools import wraps
-from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, TypeVar
 
 import datachain
 from datachain.dataset import DatasetDependency, DatasetRecord
@@ -9,7 +9,10 @@ from datachain.error import DatasetNotFoundError
 from datachain.project import Project
 
 if TYPE_CHECKING:
-    from typing_extensions import Concatenate, ParamSpec
+    from collections.abc import Callable
+    from typing import Concatenate
+
+    from typing_extensions import ParamSpec
 
     from datachain.lib.dc import DataChain
 
@@ -55,8 +58,8 @@ def _get_delta_chain(
     source_ds_project: Project,
     source_ds_version: str,
     source_ds_latest_version: str,
-    on: Union[str, Sequence[str]],
-    compare: Optional[Union[str, Sequence[str]]] = None,
+    on: str | Sequence[str],
+    compare: str | Sequence[str] | None = None,
 ) -> "DataChain":
     """Get delta chain for processing changes between versions."""
     source_dc = datachain.read_dataset(
@@ -84,11 +87,11 @@ def _get_retry_chain(
     source_ds_name: str,
     source_ds_project: Project,
     source_ds_version: str,
-    on: Union[str, Sequence[str]],
-    right_on: Optional[Union[str, Sequence[str]]],
-    delta_retry: Optional[Union[bool, str]],
+    on: str | Sequence[str],
+    right_on: str | Sequence[str] | None,
+    delta_retry: bool | str | None,
     diff_chain: "DataChain",
-) -> Optional["DataChain"]:
+) -> "DataChain | None":
     """Get retry chain for processing error records and missing records."""
     # Import here to avoid circular import
     from datachain.lib.dc import C
@@ -144,11 +147,11 @@ def _get_source_info(
     latest_version: str,
     catalog,
 ) -> tuple[
-    Optional[str],
-    Optional[Project],
-    Optional[str],
-    Optional[str],
-    Optional[list[DatasetDependency]],
+    str | None,
+    Project | None,
+    str | None,
+    str | None,
+    list[DatasetDependency] | None,
 ]:
     """Get source dataset information and dependencies.
 
@@ -190,11 +193,11 @@ def delta_retry_update(
     namespace_name: str,
     project_name: str,
     name: str,
-    on: Union[str, Sequence[str]],
-    right_on: Optional[Union[str, Sequence[str]]] = None,
-    compare: Optional[Union[str, Sequence[str]]] = None,
-    delta_retry: Optional[Union[bool, str]] = None,
-) -> tuple[Optional["DataChain"], Optional[list[DatasetDependency]], bool]:
+    on: str | Sequence[str],
+    right_on: str | Sequence[str] | None = None,
+    compare: str | Sequence[str] | None = None,
+    delta_retry: bool | str | None = None,
+) -> tuple["DataChain | None", list[DatasetDependency] | None, bool]:
     """
     Creates new chain that consists of the last version of current delta dataset
     plus diff from the source with all needed modifications.