ygg 0.1.30__py3-none-any.whl → 0.1.31__py3-none-any.whl

{ygg-0.1.30.dist-info → ygg-0.1.31.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.30
+Version: 0.1.31
 Summary: Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks
 Author: Yggdrasil contributors
 License:                                  Apache License

{ygg-0.1.30.dist-info → ygg-0.1.31.dist-info}/RECORD

@@ -1,4 +1,4 @@
-ygg-0.1.30.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+ygg-0.1.31.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
 yggdrasil/__init__.py,sha256=6OPibApplA5TF4TeixkQO_qewpaAidYX-fSDvvKYcTI,91
 yggdrasil/databricks/__init__.py,sha256=aGVve5mpoQtxSK2nfzrexjRPoutCIyaOnKZijkG4_QE,92
 yggdrasil/databricks/compute/__init__.py,sha256=TVDwPmW2SOmHmnhzZhsvrWbrxZ_lEcgqe3l9BeB-oxM,218
@@ -6,15 +6,18 @@ yggdrasil/databricks/compute/cluster.py,sha256=P4e4eBSSSykO8lzrGLLfHEZukObqcnJgC
 yggdrasil/databricks/compute/execution_context.py,sha256=ecpjx1smGYu85Gv57SUM9T1Yo3n71gyRhAVqbe7X4-0,18677
 yggdrasil/databricks/compute/remote.py,sha256=f8IJpPmEjXQh8toaS5Kt8rxkeVph7Uxrcv7utKjnT1I,1347
 yggdrasil/databricks/jobs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-yggdrasil/databricks/jobs/config.py,sha256=8Slfw4Wl7vu0kIlaUUqVqjjOgPwuULoo0rroENCbC20,11494
+yggdrasil/databricks/jobs/config.py,sha256=FI2fiiXcxps-aRRfcqqQCqnnYj_TUHQNVJMBxQBe_OU,10693
 yggdrasil/databricks/sql/__init__.py,sha256=JZpQ9eCphDf1l4yzIZ7a7OLigxqXkqOgb0Mio7Rj09A,181
 yggdrasil/databricks/sql/engine.py,sha256=edNXVdAwjfGENegoeE8cq1mxelaqsUJFtf57puZFEpY,30894
 yggdrasil/databricks/sql/exceptions.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-yggdrasil/databricks/sql/statement_result.py,sha256=o937Xx8XyXVPdnDPrYAkIjNfsyRpA6a5uzjxRpp7-1U,12909
-yggdrasil/databricks/sql/types.py,sha256=YgasSyq8sygk1h6ZOTcRwXAZWNKSuk-9g9VqlR8kJl4,5324
-yggdrasil/databricks/workspaces/__init__.py,sha256=tNNS3A_Pl9FYkQ8nGERhr4VF-hwKrvh8k1W8vTaR0uo,58
-yggdrasil/databricks/workspaces/databricks_path.py,sha256=Z0PqcJnDAwlfh4VXwfOfUhyzyTuGbZsW4mCoCBkZcCo,28284
-yggdrasil/databricks/workspaces/workspace.py,sha256=0MQAJSeJBKGHLAZyDoP5j5GLG4gjOUV4Jh2QvESOglk,19644
+yggdrasil/databricks/sql/statement_result.py,sha256=sqg3hLG-0vmMZSqQEf5QEbmb3CPvBwBKov7hPc_cPlo,12909
+yggdrasil/databricks/sql/types.py,sha256=nkoLNCan5tGRFq_rQJUz2mbz5bjr8effHvm4Gkuaf68,5661
+yggdrasil/databricks/workspaces/__init__.py,sha256=9Hk1csqJp2UiaAS5MZ7oz6KHh9JNo7A4Vo1xf61y2DM,68
+yggdrasil/databricks/workspaces/filesytem.py,sha256=e7oWaLdbhnl8evl2AjTUdhnITMem3ry2WmOFO4CQ02g,4761
+yggdrasil/databricks/workspaces/io.py,sha256=tGyjB5MWJK01lgeZC2my8mlyEEfpS9f-2MidTJNtlCI,21363
+yggdrasil/databricks/workspaces/path.py,sha256=MEPck7V0Pc92JGzgkFz5b6p8rH1m_otmewGttsAs_nQ,39091
+yggdrasil/databricks/workspaces/path_kind.py,sha256=S7f9jKLeceQNc6uZu9safASkZqHfKU-06Pq_3ixtfIU,174
+yggdrasil/databricks/workspaces/workspace.py,sha256=vVm4hD3_Bjv3qRGz6Q03LgNCaIithaJJCCAVVYTYWKM,17563
 yggdrasil/dataclasses/__init__.py,sha256=QVAvZnNl7gFYTLOVTfMkdQZf6o_WL8_UuDV1uTZ7Aeg,67
 yggdrasil/dataclasses/dataclass.py,sha256=ln-D1-bbiCLBd2khRMGs4dFoxzJEEGzHTKekWCnF2uk,5436
 yggdrasil/libs/__init__.py,sha256=88VTVKz8pnB8JSit6rlNzFm9O1ORyxqh0TQ7OBkXkQo,104
@@ -40,7 +43,7 @@ yggdrasil/types/libs.py,sha256=7-p0M4C6TnEWpUGf2nY4XshhJxtXOE_-bsYmJWU6jtk,227
 yggdrasil/types/python_arrow.py,sha256=1Ac1ZnEgmH4nLNBfrfbM6E9EfANGWtdANqlwMENQkTw,21408
 yggdrasil/types/python_defaults.py,sha256=7t0GMdgn1yJ-xhObeD3-vP091e6tdPl6qrm-wU7bbqw,8390
 yggdrasil/types/cast/__init__.py,sha256=2qiXTrUk7HRcdN8ZuDDUxSDCkqJfAfuwG0M1XA7rTzk,240
-yggdrasil/types/cast/arrow_cast.py,sha256=Eh0HTAEOSL5d_UAI0sX4qfBlTGSMkGMgDNVqd7875ec,38398
+yggdrasil/types/cast/arrow_cast.py,sha256=1brucYDBLJ9MUjQQ4ura2slIW2lUApzTc_7efq043fg,38664
 yggdrasil/types/cast/cast_options.py,sha256=R758Gfv3latG6hq_xrB4hbNPHPL51D-euwIvKEY8YYw,12149
 yggdrasil/types/cast/pandas_cast.py,sha256=-DlaZxGmisIklsZElUZI3g6nrj4R-xNDTI7q01IMB6g,7889
 yggdrasil/types/cast/polars_cast.py,sha256=Z48vvAVyRYZVYabiad22e3tn44tZ1iaSJgOOoJ1MorU,27523
@@ -49,8 +52,8 @@ yggdrasil/types/cast/registry.py,sha256=-88mq-U1pDSGbEC9PRY0zJCzloyBodXgeSRBPb6h
 yggdrasil/types/cast/spark_cast.py,sha256=IHthM78dugabGXxNNW9sSHn-olDwzXcFdIFcPo9IiXU,23021
 yggdrasil/types/cast/spark_pandas_cast.py,sha256=8PgJItF_XbyBcNuBnXkMQU3PBy3sAPEXZT9SXL2WbU4,4200
 yggdrasil/types/cast/spark_polars_cast.py,sha256=ba1UOvY1ouGCro1Np9slXmJ4TEyWnUtwVEAwxGvPLlk,8336
-ygg-0.1.30.dist-info/METADATA,sha256=BBuosBndVIX3PE-BsyjsuXVSjkHYG76cxUr-xXwD24o,19204
-ygg-0.1.30.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ygg-0.1.30.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
-ygg-0.1.30.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
-ygg-0.1.30.dist-info/RECORD,,
+ygg-0.1.31.dist-info/METADATA,sha256=6xmpGqSGuLKSCgOJkXoMJz3gHI4494h9fCT8Mx5daEQ,19204
+ygg-0.1.31.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ygg-0.1.31.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
+ygg-0.1.31.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
+ygg-0.1.31.dist-info/RECORD,,

yggdrasil/databricks/jobs/config.py

@@ -2,11 +2,11 @@ import builtins
 import dataclasses
 import datetime as dt
 import inspect
+import logging
 from dataclasses import dataclass, fields
 from enum import Enum
 from inspect import isclass
-from typing import Any, Dict, List, get_type_hints, Optional, get_origin
-import logging
+from typing import Any, Dict, List, get_type_hints, get_origin
 
 from ...libs.sparklib import SparkSession
 from ...types.cast.registry import convert
@@ -308,31 +308,3 @@ class NotebookConfig:
                 spark_session.conf.set("spark.sql.session.timeZone", "UTC")
 
         return cls.from_environment()
-
-
-class ExampleEnum(Enum):
-    """Example enum for widget demonstration"""
-    OPTION1 = "option1"
-    OPTION2 = "option2"
-    OPTION3 = "option3"
-
-
-@dataclass
-class CompleteNotebookConfig(NotebookConfig):
-    """Example JobConfig with various field types to demonstrate widget handling"""
-    # Basic types
-    text_field: str
-    integer_field: int = 42
-    float_field: float = 3.14
-    boolean_field: bool = True
-
-    # Special types
-    date_field: dt.date = dt.date(2023, 1, 1)
-    datetime_field: dt.datetime = dt.datetime(2023, 1, 1, 12, 0, 0)
-    enum_field: ExampleEnum = ExampleEnum.OPTION1
-
-    # Collection types
-    list_of_strings: List[str] = None  # Will be displayed as multiselect
-
-    # Optional fields
-    optional_text: Optional[str] = None

yggdrasil/databricks/sql/statement_result.py

@@ -96,7 +96,7 @@ class StatementResult:
         if self.is_spark_sql:
             return self._response
 
-        if not self.done and self.statement_id and time.time() - self._response_refresh_time > delay:
+        if self.statement_id and not self.done and time.time() - self._response_refresh_time > delay:
             self.response = self.workspace.sdk().statement_execution.get_statement(self.statement_id)
 
         return self._response

yggdrasil/databricks/sql/types.py

@@ -102,6 +102,18 @@ def _split_top_level_commas(s: str):
     return parts
 
 
+def _safe_bytes(obj):
+    if not isinstance(obj, bytes):
+        if not obj:
+            return b""
+
+        if not isinstance(obj, str):
+            obj = str(obj)
+
+        return obj.encode("utf-8")
+    return obj
+
+
 def parse_sql_type_to_pa(type_str: str) -> pa.DataType:
     """
     Adapted parser that:
@@ -170,6 +182,10 @@ def column_info_to_arrow_field(col: Union[SQLColumnInfo, CatalogColumnInfo]):
     if isinstance(col, CatalogColumnInfo):
         parsed = json.loads(col.type_json)
         md = parsed.get("metadata", {}) or {}
+        md = {
+            _safe_bytes(k): _safe_bytes(v)
+            for k, v in md.items()
+        }
         nullable = col.nullable
     elif isinstance(col, SQLColumnInfo):
         md = {}

yggdrasil/databricks/workspaces/__init__.py

@@ -1,2 +1,4 @@
 from .workspace import *
-from .databricks_path import *
+from .path import *
+from .io import *
+

yggdrasil/databricks/workspaces/filesytem.py

@@ -0,0 +1,161 @@
+__all__ = [
+    "DatabricksFileSystem",
+    "DatabricksFileSystemHandler"
+]
+
+from typing import TYPE_CHECKING, Any, Union, List, Optional
+
+from pyarrow import PythonFile
+from pyarrow.fs import FileSystem, FileInfo, FileSelector, PyFileSystem, FileSystemHandler
+
+if TYPE_CHECKING:
+    from ..workspaces.workspace import Workspace
+    from .path import DatabricksPath
+
+
+class DatabricksFileSystemHandler(FileSystemHandler):
+
+    def __init__(
+        self,
+        workspace: "Workspace",
+    ):
+        super().__init__()
+        self.workspace = workspace
+
+    def __enter__(self):
+        return self.connect(clone=True)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.workspace.__exit__(exc_type, exc_val, exc_tb)
+
+    def _parse_path(self, obj: Any) -> "DatabricksPath":
+        from .path import DatabricksPath
+
+        return DatabricksPath.parse(obj, workspace=self.workspace)
+
+    def connect(self, clone: bool = True):
+        workspace = self.connect(clone=clone)
+
+        if clone:
+            return DatabricksFileSystemHandler(
+                workspace=workspace
+            )
+
+        self.workspace = workspace
+        return self
+
+    def close(self):
+        self.workspace.close()
+
+    def copy_file(self, src, dest, *, chunk_size: int = 4 * 1024 * 1024):
+        src = self._parse_path(src)
+        dest = self._parse_path(dest)
+
+        with src.open("rb") as r, dest.open("wb") as w:
+            while True:
+                chunk = r.read(chunk_size)
+                if not chunk:
+                    break
+                w.write(chunk)
+
+    def create_dir(self, path, *args, recursive: bool = True, **kwargs):
+        return self._parse_path(path).mkdir(parents=recursive)
+
+    def delete_dir(self, path):
+        return self._parse_path(path).rmdir(recursive=True)
+
+    def delete_dir_contents(self, path, *args, accept_root_dir: bool = False, **kwargs):
+        return self._parse_path(path).rmdir(recursive=True)
+
+    def delete_root_dir_contents(self):
+        return self.delete_dir_contents("/", accept_root_dir=True)
+
+    def delete_file(self, path):
+        return self._parse_path(path).rmfile()
+
+    def equals(self, other: FileSystem):
+        return self == other
+
+    def from_uri(self, uri):
+        uri = self._parse_path(uri)
+
+        return self.__class__(
+            workspace=uri.workspace
+        )
+
+    def get_file_info(
+        self,
+        paths_or_selector: Union[FileSelector, str, "DatabricksPath", List[Union[str, "DatabricksPath"]]]
+    ) -> Union[FileInfo, List[FileInfo]]:
+        from .path import DatabricksPath
+
+        if isinstance(paths_or_selector, (str, DatabricksPath)):
+            result = self._parse_path(paths_or_selector).file_info
+
+            return result
+
+        if isinstance(paths_or_selector, FileSelector):
+            return self.get_file_info_selector(paths_or_selector)
+
+        return [
+            self.get_file_info(obj)
+            for obj in paths_or_selector
+        ]
+
+    def get_file_info_selector(
+        self,
+        selector: FileSelector
+    ):
+        base_dir = self._parse_path(selector.base_dir)
+
+        return [
+            p.file_info
+            for p in base_dir.ls(
+                recursive=selector.recursive,
+                allow_not_found=selector.allow_not_found
+            )
+        ]
+
+    def get_type_name(self):
+        return "dbfs"
+
+    def move(self, src, dest):
+        src = self._parse_path(src)
+
+        src.copy_to(dest)
+
+        src.remove(recursive=True)
+
+    def normalize_path(self, path):
+        return self._parse_path(path).full_path()
+
+    def open(
+        self,
+        path,
+        mode: str = "r+",
+        encoding: Optional[str] = None,
+    ):
+        return self._parse_path(path).open(mode=mode, encoding=encoding, clone=False)
+
+    def open_append_stream(self, path, compression='detect', buffer_size=None, metadata=None):
+        return self._parse_path(path).open(mode="ab")
+
+    def open_input_file(self, path, mode: str = "rb", **kwargs):
+        buf = self._parse_path(path).open(mode=mode).connect(clone=True)
+
+        return PythonFile(
+            buf,
+            mode=mode
+        )
+
+    def open_input_stream(self, path, compression='detect', buffer_size=None):
+        return self._parse_path(path).open(mode="rb")
+
+    def open_output_stream(self, path, compression='detect', buffer_size=None, metadata=None):
+        return self._parse_path(path).open(mode="wb")
+
+
+class DatabricksFileSystem(PyFileSystem):
+
+    def __init__(self, handler): # real signature unknown; restored from __doc__
+        super().__init__(handler)