speedy-utils 1.1.27__py3-none-any.whl → 1.1.29__py3-none-any.whl

speedy_utils/common/utils_io.py

@@ -1,95 +1,102 @@
 # utils/utils_io.py
 
-import bz2
-import gzip
-import io
-import json
-import lzma
-import os
-import os.path as osp
-import pickle
-import time
-import warnings
-from glob import glob
-from pathlib import Path
-from typing import IO, Any, Iterable, Optional, Union, cast
-
-from json_repair import loads as jloads
-from pydantic import BaseModel
-
+# import bz2
+# import contextlib
+# import gzip
+# import io
+# import json
+# import lzma
+# import os
+# import os.path as osp
+# import pickle
+# import time
+# import warnings
+# from collections.abc import Iterable
+# from glob import glob
+# from pathlib import Path
+# from typing import IO, Any, Optional, Union, cast
+
+
+# from pydantic import BaseModel
+
+from ..__imports import *
 from .utils_misc import mkdir_or_exist
 
 
-def dump_jsonl(list_dictionaries: list[dict], file_name: str = "output.jsonl") -> None:
+def dump_jsonl(list_dictionaries: list[dict], file_name: str = 'output.jsonl') -> None:
     """
     Dumps a list of dictionaries to a file in JSON Lines format.
     """
-    with open(file_name, "w", encoding="utf-8") as file:
+    with open(file_name, 'w', encoding='utf-8') as file:
         for dictionary in list_dictionaries:
-            file.write(json.dumps(dictionary, ensure_ascii=False) + "\n")
+            file.write(json.dumps(dictionary, ensure_ascii=False) + '\n')
 
 
-def dump_json_or_pickle(obj: Any, fname: str, ensure_ascii: bool = False, indent: int = 4) -> None:
+def dump_json_or_pickle(
+    obj: Any, fname: str, ensure_ascii: bool = False, indent: int = 4
+) -> None:
     """
     Dump an object to a file, supporting both JSON and pickle formats.
     """
     if isinstance(fname, Path):
         fname = str(fname)
     mkdir_or_exist(osp.abspath(os.path.dirname(osp.abspath(fname))))
-    if fname.endswith(".json"):
-        with open(fname, "w", encoding="utf-8") as f:
+    if fname.endswith('.json'):
+        with open(fname, 'w', encoding='utf-8') as f:
             try:
                 json.dump(obj, f, ensure_ascii=ensure_ascii, indent=indent)
             # TypeError: Object of type datetime is not JSON serializable
             except TypeError:
                 print(
-                    "Error: Object of type datetime is not JSON serializable",
+                    'Error: Object of type datetime is not JSON serializable',
                     str(obj)[:1000],
                 )
                 raise
-    elif fname.endswith(".jsonl"):
+    elif fname.endswith('.jsonl'):
         dump_jsonl(obj, fname)
-    elif fname.endswith(".pkl"):
+    elif fname.endswith('.pkl'):
         try:
-            with open(fname, "wb") as f:
+            with open(fname, 'wb') as f:
                 pickle.dump(obj, f)
         except Exception as e:
             if isinstance(obj, BaseModel):
-                data = obj.model_dump()
+                data = obj.model_dump()  # type: ignore
                 from fastcore.all import dict2obj, obj2dict
 
                 obj2 = dict2obj(data)
-                with open(fname, "wb") as f:
+                with open(fname, 'wb') as f:
                     pickle.dump(obj2, f)
             else:
-                raise ValueError(f"Error {e} while dumping {fname}") from e
+                raise ValueError(f'Error {e} while dumping {fname}') from e
 
     else:
-        raise NotImplementedError(f"File type {fname} not supported")
+        raise NotImplementedError(f'File type {fname} not supported')
 
 
 def load_json_or_pickle(fname: str, counter=0) -> Any:
     """
     Load an object from a file, supporting both JSON and pickle formats.
     """
-    if fname.endswith(".json") or fname.endswith(".jsonl"):
-        with open(fname, encoding="utf-8") as f:
+    if fname.endswith(('.json', '.jsonl')):
+        with open(fname, encoding='utf-8') as f:
             return json.load(f)
     else:
         try:
-            with open(fname, "rb") as f:
+            with open(fname, 'rb') as f:
                 return pickle.load(f)
         # EOFError: Ran out of input
         except EOFError:
             time.sleep(1)
             if counter > 5:
                 # Keep message concise and actionable
-                print(f"Corrupted cache file {fname} removed; it will be regenerated on next access")
+                print(
+                    f'Corrupted cache file {fname} removed; it will be regenerated on next access'
+                )
                 os.remove(fname)
                 raise
             return load_json_or_pickle(fname, counter + 1)
         except Exception as e:
-            raise ValueError(f"Error {e} while loading {fname}") from e
+            raise ValueError(f'Error {e} while loading {fname}') from e
 
 
 try:
@@ -104,19 +111,19 @@ except Exception:
 
 
 def fast_load_jsonl(
-    path_or_file: Union[str, os.PathLike, IO],
+    path_or_file: str | os.PathLike | IO,
     *,
     progress: bool = False,
-    desc: str = "Reading JSONL",
+    desc: str = 'Reading JSONL',
     use_orjson: bool = True,
-    encoding: str = "utf-8",
-    errors: str = "strict",
-    on_error: str = "raise",  # 'raise' | 'warn' | 'skip'
+    encoding: str = 'utf-8',
+    errors: str = 'strict',
+    on_error: str = 'raise',  # 'raise' | 'warn' | 'skip'
     skip_empty: bool = True,
-    max_lines: Optional[int] = None,
+    max_lines: int | None = None,
     use_multiworker: bool = True,
     multiworker_threshold: int = 1000000,
-    workers: Optional[int] = None,
+    workers: int | None = None,
 ) -> Iterable[Any]:
     """
     Lazily iterate objects from a JSON Lines file.
@@ -144,31 +151,52 @@ def fast_load_jsonl(
         Parsed Python objects per line.
     """
 
+    class ZstdWrapper:
+        """Context manager wrapper for zstd decompression."""
+
+        def __init__(self, path):
+            self.path = path
+            self.fh = None
+            self.stream = None
+            self.buffered = None
+
+        def __enter__(self):
+            if zstd is None:
+                raise ImportError('zstandard package required for .zst files')
+            self.fh = open(self.path, 'rb')
+            dctx = zstd.ZstdDecompressor()
+            self.stream = dctx.stream_reader(self.fh)
+            self.buffered = io.BufferedReader(self.stream)
+            return self.buffered
+
+        def __exit__(self, exc_type, exc_val, exc_tb):
+            if self.fh:
+                self.fh.close()
+
     def _open_auto(pth_or_f) -> IO[Any]:
-        if hasattr(pth_or_f, "read"):
+        if hasattr(pth_or_f, 'read'):
             # ensure binary buffer for consistent byte-length progress
             fobj = pth_or_f
             # If it's text, wrap it to binary via encoding; else just return
             if isinstance(fobj, io.TextIOBase):
                 # TextIO -> re-encode to bytes on the fly
-                return io.BufferedReader(io.BytesIO(fobj.read().encode(encoding, errors)))
+                return io.BufferedReader(
+                    io.BytesIO(fobj.read().encode(encoding, errors))
+                )
             return pth_or_f  # assume binary
         s = str(pth_or_f).lower()
-        if s.endswith(".gz"):
-            return gzip.open(pth_or_f, "rb")  # type: ignore
-        if s.endswith(".bz2"):
-            return bz2.open(pth_or_f, "rb")  # type: ignore
-        if s.endswith((".xz", ".lzma")):
-            return lzma.open(pth_or_f, "rb")  # type: ignore
-        if s.endswith((".zst", ".zstd")) and zstd is not None:
-            fh = open(pth_or_f, "rb")
-            dctx = zstd.ZstdDecompressor()
-            stream = dctx.stream_reader(fh)
-            return io.BufferedReader(stream)  # type: ignore
+        if s.endswith('.gz'):
+            return gzip.open(pth_or_f, 'rb')  # type: ignore
+        if s.endswith('.bz2'):
+            return bz2.open(pth_or_f, 'rb')  # type: ignore
+        if s.endswith(('.xz', '.lzma')):
+            return lzma.open(pth_or_f, 'rb')  # type: ignore
+        if s.endswith(('.zst', '.zstd')) and zstd is not None:
+            return ZstdWrapper(pth_or_f).__enter__()  # type: ignore
         # plain
-        return open(pth_or_f, "rb", buffering=1024 * 1024)
+        return open(pth_or_f, 'rb', buffering=1024 * 1024)
 
-    def _count_lines_fast(file_path: Union[str, os.PathLike]) -> int:
+    def _count_lines_fast(file_path: str | os.PathLike) -> int:
         """Quickly count lines in a file, handling compression."""
         try:
             f = _open_auto(file_path)
@@ -187,7 +215,7 @@ def fast_load_jsonl(
         for line_bytes in chunk_lines:
             if skip_empty and not line_bytes.strip():
                 continue
-            line_bytes = line_bytes.rstrip(b"\r\n")
+            line_bytes = line_bytes.rstrip(b'\r\n')
             try:
                 if use_orjson and orjson is not None:
                     obj = orjson.loads(line_bytes)
@@ -195,10 +223,10 @@ def fast_load_jsonl(
                     obj = json.loads(line_bytes.decode(encoding, errors))
                 results.append(obj)
             except Exception as e:
-                if on_error == "raise":
+                if on_error == 'raise':
                     raise
-                if on_error == "warn":
-                    warnings.warn(f"Skipping malformed line: {e}")
+                if on_error == 'warn':
+                    warnings.warn(f'Skipping malformed line: {e}', stacklevel=2)
                 # 'skip' and 'warn' both skip the line
                 continue
         return results
@@ -206,12 +234,12 @@ def fast_load_jsonl(
     # Check if we should use multi-worker processing
     should_use_multiworker = (
         use_multiworker
-        and not hasattr(path_or_file, "read")  # Only for file paths, not file objects
+        and not hasattr(path_or_file, 'read')  # Only for file paths, not file objects
         and max_lines is None  # Don't use multiworker if we're limiting lines
     )
 
     if should_use_multiworker:
-        line_count = _count_lines_fast(cast(Union[str, os.PathLike], path_or_file))
+        line_count = _count_lines_fast(cast(str | os.PathLike, path_or_file))
         if line_count > multiworker_threshold:
             # Use multi-worker processing
             from ..multi_worker.thread import multi_thread
@@ -236,9 +264,13 @@ def fast_load_jsonl(
 
             # Process chunks in parallel
             if progress:
-                print(f"Processing {line_count} lines with {num_workers} workers ({len(chunks)} chunks)...")
+                print(
+                    f'Processing {line_count} lines with {num_workers} workers ({len(chunks)} chunks)...'
+                )
 
-            chunk_results = multi_thread(_process_chunk, chunks, workers=num_workers, progress=progress)
+            chunk_results = multi_thread(
+                _process_chunk, chunks, workers=num_workers, progress=progress
+            )
 
             # Flatten results and yield
             if chunk_results:
@@ -257,15 +289,15 @@ def fast_load_jsonl(
         try:
             from tqdm import tqdm  # type: ignore
         except Exception as e:
-            raise ImportError("tqdm is required when progress=True") from e
+            raise ImportError('tqdm is required when progress=True') from e
         total = None
-        if not hasattr(path_or_file, "read"):
+        if not hasattr(path_or_file, 'read'):
             try:
-                path_for_size = cast(Union[str, os.PathLike], path_or_file)
+                path_for_size = cast(str | os.PathLike, path_or_file)
                 total = os.path.getsize(path_for_size)  # compressed size if any
             except Exception:
                 total = None
-        pbar = tqdm(total=total, unit="B", unit_scale=True, desc=desc)
+        pbar = tqdm(total=total, unit='B', unit_scale=True, desc=desc)
 
     line_no = 0
     try:
@@ -286,7 +318,7 @@ def fast_load_jsonl(
                     if max_lines and line_no >= max_lines:
                         break
                     continue
-                line_bytes = raw_line.rstrip(b"\r\n")
+                line_bytes = raw_line.rstrip(b'\r\n')
                 # Parse
                 try:
                     if use_orjson and orjson is not None:
@@ -294,10 +326,12 @@ def fast_load_jsonl(
                     else:
                         obj = json.loads(line_bytes.decode(encoding, errors))
                 except Exception as e:
-                    if on_error == "raise":
+                    if on_error == 'raise':
                         raise
-                    if on_error == "warn":
-                        warnings.warn(f"Skipping malformed line {line_no}: {e}")
+                    if on_error == 'warn':
+                        warnings.warn(
+                            f'Skipping malformed line {line_no}: {e}', stacklevel=2
+                        )
                     # 'skip' and 'warn' both skip the line
                     if max_lines and line_no >= max_lines:
                         break
@@ -311,10 +345,12 @@ def fast_load_jsonl(
                 try:
                     obj = json.loads(raw_line)
                 except Exception as e:
-                    if on_error == "raise":
+                    if on_error == 'raise':
                         raise
-                    if on_error == "warn":
-                        warnings.warn(f"Skipping malformed line {line_no}: {e}")
+                    if on_error == 'warn':
+                        warnings.warn(
+                            f'Skipping malformed line {line_no}: {e}', stacklevel=2
+                        )
                     if max_lines and line_no >= max_lines:
                         break
                     continue
@@ -326,14 +362,12 @@ def fast_load_jsonl(
         if pbar is not None:
             pbar.close()
         # Close only if we opened it (i.e., not an external stream)
-        if not hasattr(path_or_file, "read"):
-            try:
+        if not hasattr(path_or_file, 'read'):
+            with contextlib.suppress(Exception):
                 f.close()
-            except Exception:
-                pass
 
 
-def load_by_ext(fname: Union[str, list[str]], do_memoize: bool = False) -> Any:
+def load_by_ext(fname: str | list[str], do_memoize: bool = False) -> Any:
     """
     Load data based on file extension.
     """
@@ -346,54 +380,54 @@ def load_by_ext(fname: Union[str, list[str]], do_memoize: bool = False) -> Any:
     )
 
     try:
-        if isinstance(fname, str) and "*" in fname:
+        if isinstance(fname, str) and '*' in fname:
             paths = glob(fname)
             paths = sorted(paths)
             return multi_process(load_by_ext, paths, workers=16)
-        elif isinstance(fname, list):
+        if isinstance(fname, list):
             paths = fname
             return multi_process(load_by_ext, paths, workers=16)
 
         def load_csv(path: str, **pd_kwargs) -> Any:
             import pandas as pd
 
-            return pd.read_csv(path, engine="pyarrow", **pd_kwargs)
+            return pd.read_csv(path, engine='pyarrow', **pd_kwargs)
 
         def load_txt(path: str) -> list[str]:
-            with open(path, encoding="utf-8") as f:
+            with open(path, encoding='utf-8') as f:
                 return f.read().splitlines()
 
         def load_default(path: str) -> Any:
-            if path.endswith(".jsonl"):
+            if path.endswith('.jsonl'):
                 return list(fast_load_jsonl(path, progress=True))
-            elif path.endswith(".json"):
+            if path.endswith('.json'):
                 try:
                     return load_json_or_pickle(path)
                 except json.JSONDecodeError as exc:
-                    raise ValueError("JSON decoding failed") from exc
+                    raise ValueError('JSON decoding failed') from exc
             return load_json_or_pickle(path)
 
         handlers = {
-            ".csv": load_csv,
-            ".tsv": load_csv,
-            ".txt": load_txt,
-            ".pkl": load_default,
-            ".json": load_default,
-            ".jsonl": load_default,
+            '.csv': load_csv,
+            '.tsv': load_csv,
+            '.txt': load_txt,
+            '.pkl': load_default,
+            '.json': load_default,
+            '.jsonl': load_default,
         }
 
         ext = os.path.splitext(fname)[-1]
         load_fn = handlers.get(ext)
 
         if not load_fn:
-            raise NotImplementedError(f"File type {ext} not supported")
+            raise NotImplementedError(f'File type {ext} not supported')
 
         if do_memoize:
             load_fn = memoize(load_fn)
 
         return load_fn(fname)
     except Exception as e:
-        raise ValueError(f"Error {e} while loading {fname}") from e
+        raise ValueError(f'Error {e} while loading {fname}') from e
 
 
 def jdumps(obj, ensure_ascii=False, indent=2, **kwargs):
@@ -403,10 +437,10 @@ def jdumps(obj, ensure_ascii=False, indent=2, **kwargs):
 load_jsonl = lambda path: list(fast_load_jsonl(path))
 
 __all__ = [
-    "dump_json_or_pickle",
-    "dump_jsonl",
-    "load_by_ext",
-    "load_json_or_pickle",
-    "jdumps",
-    "jloads",
+    'dump_json_or_pickle',
+    'dump_jsonl',
+    'load_by_ext',
+    'load_json_or_pickle',
+    'jdumps',
+    'jloads',
 ]

speedy_utils/common/utils_misc.py

@@ -1,13 +1,8 @@
 # utils/utils_misc.py
+from ..__imports import *
 
-import inspect
-import os
-from collections.abc import Callable
-from typing import Any, TypeVar
 
-from pydantic import BaseModel
-
-T = TypeVar("T")
+T = TypeVar('T')
 
 
 def mkdir_or_exist(dir_name: str) -> None:
@@ -27,10 +22,10 @@ def get_arg_names(func: Callable) -> list[str]:
 
 def is_notebook() -> bool:
     try:
-        if "get_ipython" in globals().keys():
-            get_ipython = globals()["get_ipython"]
+        if 'get_ipython' in globals():
+            get_ipython = globals()['get_ipython']
             shell = get_ipython().__class__.__name__
-            if shell == "ZMQInteractiveShell":
+            if shell == 'ZMQInteractiveShell':
                 return True  # Jupyter notebook or qtconsole
         return False  # Other type (?)
     except NameError:
@@ -41,15 +36,14 @@ def convert_to_builtin_python(input_data: Any) -> Any:
     """Convert input data to built-in Python types."""
     if isinstance(input_data, dict):
         return {k: convert_to_builtin_python(v) for k, v in input_data.items()}
-    elif isinstance(input_data, list):
+    if isinstance(input_data, list):
         return [convert_to_builtin_python(v) for v in input_data]
-    elif isinstance(input_data, (int, float, str, bool, type(None))):
+    if isinstance(input_data, (int, float, str, bool, type(None))):
         return input_data
-    elif isinstance(input_data, BaseModel):
+    if isinstance(input_data, BaseModel):
         data = input_data.model_dump_json()
         return convert_to_builtin_python(data)
-    else:
-        raise ValueError(f"Unsupported type {type(input_data)}")
+    raise ValueError(f'Unsupported type {type(input_data)}')
 
 
 def dedup(items: list[T], key: Callable[[T], Any]) -> list[T]:
@@ -74,10 +68,10 @@ def dedup(items: list[T], key: Callable[[T], Any]) -> list[T]:
 
 
 __all__ = [
-    "mkdir_or_exist",
-    "flatten_list",
-    "get_arg_names",
-    "is_notebook",
-    "convert_to_builtin_python",
-    "dedup",
+    'mkdir_or_exist',
+    'flatten_list',
+    'get_arg_names',
+    'is_notebook',
+    'convert_to_builtin_python',
+    'dedup',
 ]

speedy_utils/common/utils_print.py

@@ -1,17 +1,11 @@
 # utils/utils_print.py
 
-import copy
-import pprint
-import textwrap
-from typing import Any, Union
-
-from tabulate import tabulate
-
+from ..__imports import *
 from .notebook_utils import display_pretty_table_html
 
 
 # Flattening the dictionary using "." notation for keys
-def flatten_dict(d, parent_key="", sep="."):
+def flatten_dict(d, parent_key='', sep='.'):
     items = []
     for k, v in d.items():
         new_key = parent_key + sep + k if parent_key else k
@@ -24,22 +18,22 @@ def flatten_dict(d, parent_key="", sep="."):
 
 def fprint(
     input_data: Any,
-    key_ignore: Union[list[str], None] = None,
-    key_keep: Union[list[str], None] = None,
+    key_ignore: list[str] | None = None,
+    key_keep: list[str] | None = None,
     max_width: int = 100,
     indent: int = 2,
-    depth: Union[int, None] = None,
-    table_format: str = "grid",
+    depth: int | None = None,
+    table_format: str = 'grid',
     str_wrap_width: int = 80,
     grep=None,
     is_notebook=None,
     f=print,
-) -> Union[None, str]:
+) -> None | str:
     """
     Pretty print structured data.
     """
     if isinstance(input_data, list):
-        for i, item in enumerate(input_data):
+        for _i, item in enumerate(input_data):
             fprint(
                 item,
                 key_ignore,
@@ -53,7 +47,7 @@ def fprint(
                 is_notebook,
                 f,
             )
-            print("\n" + "-" * 100 + "\n")
+            print('\n' + '-' * 100 + '\n')
 
     from speedy_utils import is_notebook as is_interactive
 
@@ -61,24 +55,24 @@ def fprint(
     if is_notebook is None:
         is_notebook = is_interactive()
     if isinstance(input_data, list):
-        if all(hasattr(item, "toDict") for item in input_data):
+        if all(hasattr(item, 'toDict') for item in input_data):
             input_data = [item.toDict() for item in input_data]
-    elif hasattr(input_data, "toDict"):
+    elif hasattr(input_data, 'toDict'):
         input_data = input_data.toDict()
 
     if isinstance(input_data, list):
-        if all(hasattr(item, "to_dict") for item in input_data):
+        if all(hasattr(item, 'to_dict') for item in input_data):
             input_data = [item.to_dict() for item in input_data]
-    elif hasattr(input_data, "to_dict"):
+    elif hasattr(input_data, 'to_dict'):
         input_data = input_data.to_dict()
 
     if isinstance(input_data, list):
-        if all(hasattr(item, "model_dump") for item in input_data):
+        if all(hasattr(item, 'model_dump') for item in input_data):
             input_data = [item.model_dump() for item in input_data]
-    elif hasattr(input_data, "model_dump"):
+    elif hasattr(input_data, 'model_dump'):
         input_data = input_data.model_dump()
     if not isinstance(input_data, (dict, str)):
-        raise ValueError("Input data must be a dictionary or string")
+        raise ValueError('Input data must be a dictionary or string')
 
     if isinstance(input_data, dict):
         input_data = flatten_dict(input_data)
@@ -89,7 +83,7 @@ def fprint(
     def remove_keys(d: dict, keys: list[str]) -> dict:
         """Remove specified keys from a dictionary."""
         for key in keys:
-            parts = key.split(".")
+            parts = key.split('.')
             sub_dict = d
             for part in parts[:-1]:
                 sub_dict = sub_dict.get(part, {})
@@ -100,7 +94,7 @@ def fprint(
         """Keep only specified keys in a dictionary."""
         result = {}
         for key in keys:
-            parts = key.split(".")
+            parts = key.split('.')
             sub_source = d
             sub_result = result
             for part in parts[:-1]:
@@ -112,7 +106,7 @@ def fprint(
                 sub_result[parts[-1]] = copy.deepcopy(sub_source.get(parts[-1]))
         return result
 
-    if hasattr(input_data, "to_dict") and not isinstance(input_data, str):
+    if hasattr(input_data, 'to_dict') and not isinstance(input_data, str):
         input_data = input_data.to_dict()
 
     processed_data = copy.deepcopy(input_data)
@@ -132,7 +126,7 @@ def fprint(
         f(
             tabulate(
                 table,
-                headers=["Key", "Value"],
+                headers=['Key', 'Value'],
                 tablefmt=table_format,
                 maxcolwidths=[None, max_width],
             )
@@ -148,6 +142,6 @@ def fprint(
 
 
 __all__ = [
-    "flatten_dict",
-    "fprint",
+    'flatten_dict',
+    'fprint',
 ]