mlrun 1.7.2rc3__py3-none-any.whl → 1.8.0rc2__py3-none-any.whl

mlrun/datastore/alibaba_oss.py

@@ -15,6 +15,7 @@
 import time
 from datetime import datetime
 from pathlib import Path
+from typing import Optional
 from urllib.parse import urlparse
 
 import oss2
@@ -28,7 +29,9 @@ from .base import DataStore, FileStats, make_datastore_schema_sanitizer
 class OSSStore(DataStore):
     using_bucket = True
 
-    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+    def __init__(
+        self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
+    ):
         super().__init__(parent, name, schema, endpoint, secrets)
         # will be used in case user asks to assume a role and work through fsspec
 

mlrun/datastore/azure_blob.py

@@ -14,6 +14,7 @@
 
 import time
 from pathlib import Path
+from typing import Optional
 from urllib.parse import urlparse
 
 from azure.storage.blob import BlobServiceClient
@@ -36,7 +37,9 @@ class AzureBlobStore(DataStore):
         1024 * 1024 * 8
     )  # for service_client property only, does not affect filesystem
 
-    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+    def __init__(
+        self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
+    ):
         super().__init__(parent, name, schema, endpoint, secrets=secrets)
         self._service_client = None
         self._storage_options = None

mlrun/datastore/base.py

@@ -48,7 +48,7 @@ class FileStats:
 class DataStore:
     using_bucket = False
 
-    def __init__(self, parent, name, kind, endpoint="", secrets: dict = None):
+    def __init__(self, parent, name, kind, endpoint="", secrets: Optional[dict] = None):
         self._parent = parent
         self.kind = kind
         self.name = name
@@ -500,12 +500,18 @@ class DataItem:
         """DataItem url e.g. /dir/path, s3://bucket/path"""
         return self._url
 
-    def get(self, size=None, offset=0, encoding=None):
+    def get(
+        self,
+        size: Optional[int] = None,
+        offset: int = 0,
+        encoding: Optional[str] = None,
+    ) -> Union[bytes, str]:
         """read all or a byte range and return the content
 
         :param size:     number of bytes to get
         :param offset:   fetch from offset (in bytes)
         :param encoding: encoding (e.g. "utf-8") for converting bytes to str
+        :return:         the bytes/str content
         """
         body = self._store.get(self._path, size=size, offset=offset)
         if encoding and isinstance(body, bytes):
@@ -519,7 +525,7 @@ class DataItem:
         """
         self._store.download(self._path, target_path)
 
-    def put(self, data, append=False):
+    def put(self, data: Union[bytes, str], append: bool = False) -> None:
         """write/upload the data, append is only supported by some datastores
 
         :param data:   data (bytes/str) to write
@@ -687,7 +693,9 @@ def basic_auth_header(user, password):
 
 
 class HttpStore(DataStore):
-    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+    def __init__(
+        self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
+    ):
         super().__init__(parent, name, schema, endpoint, secrets)
         self._https_auth_token = None
         self._schema = schema

mlrun/datastore/datastore.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Optional
 from urllib.parse import urlparse
 
 from mergedeep import merge
@@ -178,12 +179,17 @@ class StoreManager:
         # which accepts a feature vector uri and generate the offline vector (parquet) for it if it doesnt exist
         if not target and not allow_empty_resources:
             raise mlrun.errors.MLRunInvalidArgumentError(
-                f"resource {url} does not have a valid/persistent offline target"
+                f"Resource {url} does not have a valid/persistent offline target"
             )
         return resource, target or ""
 
     def object(
-        self, url, key="", project="", allow_empty_resources=None, secrets: dict = None
+        self,
+        url,
+        key="",
+        project="",
+        allow_empty_resources=None,
+        secrets: Optional[dict] = None,
     ) -> DataItem:
         meta = artifact_url = None
         if is_store_uri(url):
@@ -205,7 +211,7 @@ class StoreManager:
         )
 
     def get_or_create_store(
-        self, url, secrets: dict = None, project_name=""
+        self, url, secrets: Optional[dict] = None, project_name=""
     ) -> (DataStore, str, str):
         schema, endpoint, parsed_url = parse_url(url)
         subpath = parsed_url.path

mlrun/datastore/datastore_profile.py

@@ -19,7 +19,7 @@ import typing
 import warnings
 from urllib.parse import ParseResult, urlparse, urlunparse
 
-import pydantic
+import pydantic.v1
 from mergedeep import merge
 
 import mlrun
@@ -28,15 +28,15 @@ import mlrun.errors
 from ..secrets import get_secret_or_env
 
 
-class DatastoreProfile(pydantic.BaseModel):
+class DatastoreProfile(pydantic.v1.BaseModel):
     type: str
     name: str
     _private_attributes: list = ()
 
     class Config:
-        extra = pydantic.Extra.forbid
+        extra = pydantic.v1.Extra.forbid
 
-    @pydantic.validator("name")
+    @pydantic.v1.validator("name")
     @classmethod
     def lower_case(cls, v):
         return v.lower()
@@ -75,14 +75,14 @@ class TemporaryClientDatastoreProfiles(metaclass=mlrun.utils.singleton.Singleton
 
 
 class DatastoreProfileBasic(DatastoreProfile):
-    type: str = pydantic.Field("basic")
+    type: str = pydantic.v1.Field("basic")
     _private_attributes = "private"
     public: str
     private: typing.Optional[str] = None
 
 
 class DatastoreProfileKafkaTarget(DatastoreProfile):
-    type: str = pydantic.Field("kafka_target")
+    type: str = pydantic.v1.Field("kafka_target")
     _private_attributes = "kwargs_private"
     bootstrap_servers: typing.Optional[str] = None
     brokers: typing.Optional[str] = None
@@ -123,7 +123,7 @@ class DatastoreProfileKafkaTarget(DatastoreProfile):
 
 
 class DatastoreProfileKafkaSource(DatastoreProfile):
-    type: str = pydantic.Field("kafka_source")
+    type: str = pydantic.v1.Field("kafka_source")
     _private_attributes = ("kwargs_private", "sasl_user", "sasl_pass")
     brokers: typing.Union[str, list[str]]
     topics: typing.Union[str, list[str]]
@@ -162,7 +162,7 @@ class DatastoreProfileKafkaSource(DatastoreProfile):
 
 
 class DatastoreProfileV3io(DatastoreProfile):
-    type: str = pydantic.Field("v3io")
+    type: str = pydantic.v1.Field("v3io")
     v3io_access_key: typing.Optional[str] = None
     _private_attributes = "v3io_access_key"
 
@@ -178,7 +178,7 @@ class DatastoreProfileV3io(DatastoreProfile):
 
 
 class DatastoreProfileS3(DatastoreProfile):
-    type: str = pydantic.Field("s3")
+    type: str = pydantic.v1.Field("s3")
     _private_attributes = ("access_key_id", "secret_key")
     endpoint_url: typing.Optional[str] = None
     force_non_anonymous: typing.Optional[str] = None
@@ -188,7 +188,7 @@ class DatastoreProfileS3(DatastoreProfile):
     secret_key: typing.Optional[str] = None
     bucket: typing.Optional[str] = None
 
-    @pydantic.validator("bucket")
+    @pydantic.v1.validator("bucket")
     @classmethod
     def check_bucket(cls, v):
         if not v:
@@ -226,7 +226,7 @@ class DatastoreProfileS3(DatastoreProfile):
 
 
 class DatastoreProfileRedis(DatastoreProfile):
-    type: str = pydantic.Field("redis")
+    type: str = pydantic.v1.Field("redis")
     _private_attributes = ("username", "password")
     endpoint_url: str
     username: typing.Optional[str] = None
@@ -269,7 +269,7 @@ class DatastoreProfileRedis(DatastoreProfile):
 
 
 class DatastoreProfileDBFS(DatastoreProfile):
-    type: str = pydantic.Field("dbfs")
+    type: str = pydantic.v1.Field("dbfs")
     _private_attributes = ("token",)
     endpoint_url: typing.Optional[str] = None  # host
     token: typing.Optional[str] = None
@@ -287,13 +287,13 @@ class DatastoreProfileDBFS(DatastoreProfile):
 
 
 class DatastoreProfileGCS(DatastoreProfile):
-    type: str = pydantic.Field("gcs")
+    type: str = pydantic.v1.Field("gcs")
     _private_attributes = ("gcp_credentials",)
     credentials_path: typing.Optional[str] = None  # path to file.
     gcp_credentials: typing.Optional[typing.Union[str, dict]] = None
     bucket: typing.Optional[str] = None
 
-    @pydantic.validator("bucket")
+    @pydantic.v1.validator("bucket")
     @classmethod
     def check_bucket(cls, v):
         if not v:
@@ -304,7 +304,7 @@ class DatastoreProfileGCS(DatastoreProfile):
             )
         return v
 
-    @pydantic.validator("gcp_credentials", pre=True, always=True)
+    @pydantic.v1.validator("gcp_credentials", pre=True, always=True)
     @classmethod
     def convert_dict_to_json(cls, v):
         if isinstance(v, dict):
@@ -332,7 +332,7 @@ class DatastoreProfileGCS(DatastoreProfile):
 
 
 class DatastoreProfileAzureBlob(DatastoreProfile):
-    type: str = pydantic.Field("az")
+    type: str = pydantic.v1.Field("az")
     _private_attributes = (
         "connection_string",
         "account_key",
@@ -350,7 +350,7 @@ class DatastoreProfileAzureBlob(DatastoreProfile):
     credential: typing.Optional[str] = None
     container: typing.Optional[str] = None
 
-    @pydantic.validator("container")
+    @pydantic.v1.validator("container")
     @classmethod
     def check_container(cls, v):
         if not v:
@@ -392,7 +392,7 @@ class DatastoreProfileAzureBlob(DatastoreProfile):
 
 
 class DatastoreProfileHdfs(DatastoreProfile):
-    type: str = pydantic.Field("hdfs")
+    type: str = pydantic.v1.Field("hdfs")
     _private_attributes = "token"
     host: typing.Optional[str] = None
     port: typing.Optional[int] = None
@@ -415,7 +415,7 @@ class DatastoreProfileHdfs(DatastoreProfile):
         return f"webhdfs://{self.host}:{self.http_port}{subpath}"
 
 
-class DatastoreProfile2Json(pydantic.BaseModel):
+class DatastoreProfile2Json(pydantic.v1.BaseModel):
     @staticmethod
     def _to_json(attributes):
         # First, base64 encode the values
@@ -489,7 +489,7 @@ class DatastoreProfile2Json(pydantic.BaseModel):
             )
 
 
-def datastore_profile_read(url, project_name="", secrets: dict = None):
+def datastore_profile_read(url, project_name="", secrets: typing.Optional[dict] = None):
     parsed_url = urlparse(url)
     if parsed_url.scheme.lower() != "ds":
         raise mlrun.errors.MLRunInvalidArgumentError(

mlrun/datastore/dbfs_store.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import pathlib
+from typing import Optional
 
 from fsspec.implementations.dbfs import DatabricksFile, DatabricksFileSystem
 from fsspec.registry import get_filesystem_class
@@ -81,7 +82,9 @@ class DatabricksFileSystemDisableCache(DatabricksFileSystem):
 
 # dbfs objects will be represented with the following URL: dbfs://<path>
 class DBFSStore(DataStore):
-    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+    def __init__(
+        self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
+    ):
         super().__init__(parent, name, schema, endpoint, secrets=secrets)
 
     @property

mlrun/datastore/filestore.py

@@ -14,6 +14,7 @@
 import time
 from os import listdir, makedirs, path, stat
 from shutil import copyfile
+from typing import Optional
 
 import fsspec
 
@@ -23,7 +24,9 @@ from .base import DataStore, FileStats
 
 
 class FileStore(DataStore):
-    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+    def __init__(
+        self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
+    ):
         super().__init__(parent, name, "file", endpoint, secrets=secrets)
 
         self._item_path, self._real_path = None, None

mlrun/datastore/google_cloud_storage.py

@@ -14,6 +14,7 @@
 import json
 import os
 from pathlib import Path
+from typing import Optional
 
 from fsspec.registry import get_filesystem_class
 from google.auth.credentials import Credentials
@@ -33,7 +34,9 @@ class GoogleCloudStorageStore(DataStore):
     workers = 8
     chunk_size = 32 * 1024 * 1024
 
-    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+    def __init__(
+        self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
+    ):
         super().__init__(parent, name, schema, endpoint, secrets=secrets)
         self._storage_client = None
         self._storage_options = None

mlrun/datastore/hdfs.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from typing import Optional
 from urllib.parse import urlparse
 
 import fsspec
@@ -20,7 +21,9 @@ from mlrun.datastore.base import DataStore
 
 
 class HdfsStore(DataStore):
-    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+    def __init__(
+        self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
+    ):
         super().__init__(parent, name, schema, endpoint, secrets)
 
         self.host = self._get_secret_or_env("HDFS_HOST")

mlrun/datastore/inmem.py

@@ -17,6 +17,7 @@ from io import BytesIO, StringIO
 import pandas as pd
 
 import mlrun
+import mlrun.utils.helpers
 
 from .base import DataStore, FileStats
 
@@ -35,7 +36,9 @@ class InMemoryStore(DataStore):
 
     def _get_item(self, key):
         if key not in self._items:
-            raise ValueError(f"item {key} not found in memory store")
+            raise mlrun.errors.MLRunNotFoundError(
+                f"item {key} not found in memory store"
+            )
         return self._items[key]
 
     def get(self, key, size=None, offset=0):

mlrun/datastore/redis.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
 from urllib.parse import urlparse
 
 import redis
@@ -30,7 +31,9 @@ class RedisStore(DataStore):
     - key and value sizes are limited to 512MB
     """
 
-    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+    def __init__(
+        self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
+    ):
         redis_default_port = "6379"
         super().__init__(parent, name, schema, endpoint, secrets=secrets)
         self.headers = None

mlrun/datastore/s3.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import time
+from typing import Optional
 
 import boto3
 from boto3.s3.transfer import TransferConfig
@@ -26,7 +27,9 @@ from .base import DataStore, FileStats, get_range, make_datastore_schema_sanitiz
 class S3Store(DataStore):
     using_bucket = True
 
-    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+    def __init__(
+        self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
+    ):
         super().__init__(parent, name, schema, endpoint, secrets)
         # will be used in case user asks to assume a role and work through fsspec
         self._temp_credentials = None

mlrun/datastore/sources.py

@@ -181,10 +181,10 @@ class CSVSource(BaseSourceDriver):
     def __init__(
         self,
         name: str = "",
-        path: str = None,
-        attributes: dict[str, object] = None,
-        key_field: str = None,
-        schedule: str = None,
+        path: Optional[str] = None,
+        attributes: Optional[dict[str, object]] = None,
+        key_field: Optional[str] = None,
+        schedule: Optional[str] = None,
         parse_dates: Union[None, int, str, list[int], list[str]] = None,
         **kwargs,
     ):
@@ -308,11 +308,11 @@ class ParquetSource(BaseSourceDriver):
     def __init__(
         self,
         name: str = "",
-        path: str = None,
-        attributes: dict[str, object] = None,
-        key_field: str = None,
-        time_field: str = None,
-        schedule: str = None,
+        path: Optional[str] = None,
+        attributes: Optional[dict[str, object]] = None,
+        key_field: Optional[str] = None,
+        time_field: Optional[str] = None,
+        schedule: Optional[str] = None,
         start_time: Optional[Union[datetime, str]] = None,
         end_time: Optional[Union[datetime, str]] = None,
         additional_filters: Optional[list[Union[tuple, list]]] = None,
@@ -392,7 +392,9 @@ class ParquetSource(BaseSourceDriver):
         )
 
     @classmethod
-    def from_dict(cls, struct=None, fields=None, deprecated_fields: dict = None):
+    def from_dict(
+        cls, struct=None, fields=None, deprecated_fields: Optional[dict] = None
+    ):
         new_obj = super().from_dict(
             struct=struct, fields=fields, deprecated_fields=deprecated_fields
         )
@@ -564,18 +566,18 @@ class BigQuerySource(BaseSourceDriver):
     def __init__(
         self,
         name: str = "",
-        table: str = None,
-        max_results_for_table: int = None,
-        query: str = None,
-        materialization_dataset: str = None,
-        chunksize: int = None,
-        key_field: str = None,
-        time_field: str = None,
-        schedule: str = None,
+        table: Optional[str] = None,
+        max_results_for_table: Optional[int] = None,
+        query: Optional[str] = None,
+        materialization_dataset: Optional[str] = None,
+        chunksize: Optional[int] = None,
+        key_field: Optional[str] = None,
+        time_field: Optional[str] = None,
+        schedule: Optional[str] = None,
         start_time=None,
         end_time=None,
-        gcp_project: str = None,
-        spark_options: dict = None,
+        gcp_project: Optional[str] = None,
+        spark_options: Optional[dict] = None,
         **kwargs,
     ):
         if query and table:
@@ -776,19 +778,19 @@ class SnowflakeSource(BaseSourceDriver):
     def __init__(
         self,
         name: str = "",
-        key_field: str = None,
-        attributes: dict[str, object] = None,
-        time_field: str = None,
-        schedule: str = None,
+        key_field: Optional[str] = None,
+        attributes: Optional[dict[str, object]] = None,
+        time_field: Optional[str] = None,
+        schedule: Optional[str] = None,
         start_time=None,
         end_time=None,
-        query: str = None,
-        url: str = None,
-        user: str = None,
-        database: str = None,
-        schema: str = None,
-        db_schema: str = None,
-        warehouse: str = None,
+        query: Optional[str] = None,
+        url: Optional[str] = None,
+        user: Optional[str] = None,
+        database: Optional[str] = None,
+        schema: Optional[str] = None,
+        db_schema: Optional[str] = None,
+        warehouse: Optional[str] = None,
         **kwargs,
     ):
         # TODO: Remove in 1.9.0
@@ -850,9 +852,9 @@ class CustomSource(BaseSourceDriver):
 
     def __init__(
         self,
-        class_name: str = None,
+        class_name: Optional[str] = None,
         name: str = "",
-        schedule: str = None,
+        schedule: Optional[str] = None,
         **attributes,
     ):
         attributes = attributes or {}
@@ -930,12 +932,12 @@ class OnlineSource(BaseSourceDriver):
 
     def __init__(
         self,
-        name: str = None,
-        path: str = None,
-        attributes: dict[str, object] = None,
-        key_field: str = None,
-        time_field: str = None,
-        workers: int = None,
+        name: Optional[str] = None,
+        path: Optional[str] = None,
+        attributes: Optional[dict[str, object]] = None,
+        key_field: Optional[str] = None,
+        time_field: Optional[str] = None,
+        workers: Optional[int] = None,
     ):
         super().__init__(name, path, attributes, key_field, time_field)
         self.online = True
@@ -986,7 +988,7 @@ class StreamSource(OnlineSource):
         seek_to="earliest",
         shards=1,
         retention_in_hours=24,
-        extra_attributes: dict = None,
+        extra_attributes: Optional[dict] = None,
         **kwargs,
     ):
         """
@@ -1168,7 +1170,7 @@ class KafkaSource(OnlineSource):
         self,
         num_partitions: int = 4,
         replication_factor: int = 1,
-        topics: list[str] = None,
+        topics: Optional[list[str]] = None,
     ):
         """
         Create Kafka topics with the specified number of partitions and replication factor.
@@ -1226,16 +1228,16 @@ class SQLSource(BaseSourceDriver):
     def __init__(
         self,
         name: str = "",
-        chunksize: int = None,
-        key_field: str = None,
-        time_field: str = None,
-        schedule: str = None,
+        chunksize: Optional[int] = None,
+        key_field: Optional[str] = None,
+        time_field: Optional[str] = None,
+        schedule: Optional[str] = None,
         start_time: Optional[Union[datetime, str]] = None,
         end_time: Optional[Union[datetime, str]] = None,
-        db_url: str = None,
-        table_name: str = None,
-        spark_options: dict = None,
-        parse_dates: list[str] = None,
+        db_url: Optional[str] = None,
+        table_name: Optional[str] = None,
+        spark_options: Optional[dict] = None,
+        parse_dates: Optional[list[str]] = None,
         **kwargs,
     ):
         """

mlrun/datastore/store_resources.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
-
 import mlrun
 import mlrun.artifacts
 from mlrun.config import config