dycw-utilities 0.109.13__py3-none-any.whl → 0.109.15__py3-none-any.whl

{dycw_utilities-0.109.13.dist-info → dycw_utilities-0.109.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dycw-utilities
-Version: 0.109.13
+Version: 0.109.15
 Author-email: Derek Wan <d.wan@icloud.com>
 License-File: LICENSE
 Requires-Python: >=3.12

{dycw_utilities-0.109.13.dist-info → dycw_utilities-0.109.15.dist-info}/RECORD

@@ -1,5 +1,5 @@
-utilities/__init__.py,sha256=8b_ivOe7dAl9N0yVvE9__TwO49X5ejp8NxQ28w6Xp0g,61
-utilities/altair.py,sha256=5WPwsHJoM-ZDh9iek4eSqf8_Ifb3w3KDwi7f2mrtXn4,9174
+utilities/__init__.py,sha256=qZyUzwQ2zZu6ZkKszEeJGs10nbVWvqm_izEvqG8MgXM,61
+utilities/altair.py,sha256=Gpja-flOo-Db0PIPJLJsgzAlXWoKUjPU1qY-DQ829ek,9156
 utilities/astor.py,sha256=xuDUkjq0-b6fhtwjhbnebzbqQZAjMSHR1IIS5uOodVg,777
 utilities/asyncio.py,sha256=41oQUurWMvadFK5gFnaG21hMM0Vmfn2WS6OpC0R9mas,14757
 utilities/atomicwrites.py,sha256=geFjn9Pwn-tTrtoGjDDxWli9NqbYfy3gGL6ZBctiqSo,5393
@@ -26,8 +26,9 @@ utilities/hashlib.py,sha256=SVTgtguur0P4elppvzOBbLEjVM3Pea0eWB61yg2ilxo,309
 utilities/http.py,sha256=WcahTcKYRtZ04WXQoWt5EGCgFPcyHD3EJdlMfxvDt-0,946
 utilities/hypothesis.py,sha256=sLqYcrFn0I3o0R7maqliIERRtAcREk2CKG8rSHY1t5U,46205
 utilities/ipython.py,sha256=V2oMYHvEKvlNBzxDXdLvKi48oUq2SclRg5xasjaXStw,763
-utilities/iterables.py,sha256=GRuwMOPP7jqOnDdF113XYPjVg23VYn0ts1UlV-1Tzzo,45014
+utilities/iterables.py,sha256=2Yy9gZ7BR4LXR4nlX7outFAjd4dpb3lgUo7ji_sdylY,45076
 utilities/jupyter.py,sha256=ft5JA7fBxXKzP-L9W8f2-wbF0QeYc_2uLQNFDVk4Z-M,2917
+utilities/lightweight_charts.py,sha256=0xNfcsrgFI0R9xL25LtSm-W5yhfBI93qQNT6HyaXAhg,2769
 utilities/logging.py,sha256=opIwFjGKOYyMntVeCsFNXOmTY2z02hMf2UtCB76SaI4,25142
 utilities/loguru.py,sha256=MEMQVWrdECxk1e3FxGzmOf21vWT9j8CAir98SEXFKPA,3809
 utilities/luigi.py,sha256=fpH9MbxJDuo6-k9iCXRayFRtiVbUtibCJKugf7ygpv0,5988
@@ -45,8 +46,8 @@ utilities/pathlib.py,sha256=31WPMXdLIyXgYOMMl_HOI2wlo66MGSE-cgeelk-Lias,1410
 utilities/period.py,sha256=ikHXsWtDLr553cfH6p9mMaiCnIAP69B7q84ckWV3HaA,10884
 utilities/pickle.py,sha256=Bhvd7cZl-zQKQDFjUerqGuSKlHvnW1K2QXeU5UZibtg,657
 utilities/platform.py,sha256=NU7ycTvAXAG-fdYmDXaM1m4EOml2cGiaYwaUzfzSqyU,1767
-utilities/polars.py,sha256=nB2pfK8N8HRpPE_tdbiTfFGLWC_TekAqgHlYDhnUzAM,52169
-utilities/polars_ols.py,sha256=AQe3RFOMv8CEI_ZCoscb_-PxB4JWjO0TAEmk8DKLeaI,2138
+utilities/polars.py,sha256=aOQNVyV04qYZjg7Exi6zYERhSQoCMzBP74oufxqANFY,52167
+utilities/polars_ols.py,sha256=Uc9V5kvlWZ5cU93lKZ-cfAKdVFFw81tqwLW9PxtUvMs,5618
 utilities/pqdm.py,sha256=foRytQybmOQ05pjt5LF7ANyzrIa--4ScDE3T2wd31a4,3118
 utilities/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 utilities/pydantic.py,sha256=f6qtR5mO2YMuyvNmbaEj5YeD9eGA4YYfb7Bjzh9jUs0,1845
@@ -85,7 +86,7 @@ utilities/warnings.py,sha256=yUgjnmkCRf6QhdyAXzl7u0qQFejhQG3PrjoSwxpbHrs,1819
 utilities/whenever.py,sha256=TjoTAJ1R27-rKXiXzdE4GzPidmYqm0W58XydDXp-QZM,17786
 utilities/zipfile.py,sha256=24lQc9ATcJxHXBPc_tBDiJk48pWyRrlxO2fIsFxU0A8,699
 utilities/zoneinfo.py,sha256=-DQz5a0Ikw9jfSZtL0BEQkXOMC9yGn_xiJYNCLMiqEc,1989
-dycw_utilities-0.109.13.dist-info/METADATA,sha256=7LZJJAQIYMDDVTGd220fgWLkFlrPCQQYDATGv4AnkeI,13005
-dycw_utilities-0.109.13.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dycw_utilities-0.109.13.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
-dycw_utilities-0.109.13.dist-info/RECORD,,
+dycw_utilities-0.109.15.dist-info/METADATA,sha256=sqk1b4cCCM8V737VBkLkTToWlAJ5Mb7Co4LmlynJ4Tg,13005
+dycw_utilities-0.109.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dycw_utilities-0.109.15.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
+dycw_utilities-0.109.15.dist-info/RECORD,,

utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.109.13"
+__version__ = "0.109.15"

utilities/altair.py

@@ -22,7 +22,6 @@ from altair import (
     vconcat,
 )
 from altair.utils.schemapi import Undefined
-from polars import Date, Datetime
 
 from utilities.functions import ensure_bytes, ensure_number
 from utilities.iterables import always_iterable
@@ -66,7 +65,7 @@ def plot_dataframes(
 ) -> VConcatChart:
     """Plot a DataFrame as a set of time series, with a multi-line tooltip."""
     import polars as pl
-    from polars import int_range
+    from polars import Date, Datetime, int_range
 
     from utilities.polars import replace_time_zone
 

utilities/iterables.py

@@ -1009,7 +1009,7 @@ def one(*iterables: Iterable[_T]) -> _T:
     try:
         first = next(it)
     except StopIteration:
-        raise OneEmptyError from None
+        raise OneEmptyError(iterables=iterables) from None
     try:
         second = next(it)
     except StopIteration:
@@ -1018,19 +1018,19 @@ def one(*iterables: Iterable[_T]) -> _T:
 
 
 @dataclass(kw_only=True, slots=True)
-class OneError(Exception): ...
+class OneError(Exception, Generic[_T]):
+    iterables: tuple[Iterable[_T], ...]
 
 
 @dataclass(kw_only=True, slots=True)
-class OneEmptyError(OneError):
+class OneEmptyError(OneError[_T]):
     @override
     def __str__(self) -> str:
-        return "Iterable(s) must not be empty"
+        return f"Iterable(s) {get_repr(self.iterables)} must not be empty"
 
 
 @dataclass(kw_only=True, slots=True)
 class OneNonUniqueError(OneError, Generic[_T]):
-    iterables: tuple[Iterable[_T], ...]
     first: _T
     second: _T
 

utilities/lightweight_charts.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, override
+
+from utilities.iterables import OneEmptyError, OneNonUniqueError, one
+from utilities.reprlib import get_repr
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from lightweight_charts import AbstractChart, Chart
+    from lightweight_charts.abstract import SeriesCommon
+    from polars import DataFrame
+    from polars._typing import SchemaDict
+
+    from utilities.types import PathLike
+
+
+##
+
+
+def save_chart(chart: Chart, path: PathLike, /, *, overwrite: bool = False) -> None:
+    """Atomically save a chart to disk."""
+    from utilities.atomicwrites import writer  # pragma: no cover
+
+    chart.show(block=False)  # pragma: no cover
+    with (  # pragma: no cover
+        writer(path, overwrite=overwrite) as temp,
+        temp.open(mode="wb") as fh,
+    ):
+        _ = fh.write(chart.screenshot())
+    chart.exit()  # pragma: no cover
+
+
+##
+
+
+def set_dataframe(df: DataFrame, obj: AbstractChart | SeriesCommon, /) -> None:
+    """Set a `polars` DataFrame onto a Chart."""
+    from polars import Date, Datetime, col  # pragma: no cover
+
+    try:
+        name = one(k for k, v in df.schema.items() if isinstance(v, Date | Datetime))
+    except OneEmptyError:
+        raise _SetDataFrameEmptyError(schema=df.schema) from None
+    except OneNonUniqueError as error:
+        raise _SetDataFrameNonUniqueError(
+            schema=df.schema, first=error.first, second=error.second
+        ) from None
+    return obj.set(
+        df.select(
+            col(name).alias("date").dt.strftime("iso"),
+            *[c for c in df.columns if c != name],
+        ).to_pandas()
+    )
+
+
+@dataclass(kw_only=True, slots=True)
+class SetDataFrameError(Exception):
+    schema: SchemaDict
+
+
+@dataclass(kw_only=True, slots=True)
+class _SetDataFrameEmptyError(SetDataFrameError):
+    @override
+    def __str__(self) -> str:
+        return "At least 1 column must be of date/datetime type; got 0"
+
+
+@dataclass(kw_only=True, slots=True)
+class _SetDataFrameNonUniqueError(SetDataFrameError):
+    first: str
+    second: str
+
+    @override
+    def __str__(self) -> str:
+        return f"{get_repr(self.schema)} must contain exactly 1 date/datetime column; got {self.first!r}, {self.second!r} and perhaps more"
+
+
+##
+
+
+@asynccontextmanager
+async def yield_chart(chart: Chart, /) -> AsyncIterator[None]:
+    """Yield a chart for visualization in a notebook."""
+    try:  # pragma: no cover
+        yield await chart.show_async()
+    except BaseException:  # pragma: no cover  # noqa: BLE001, S110
+        pass
+    finally:  # pragma: no cover
+        chart.exit()
+
+
+__all__ = ["save_chart", "set_dataframe", "yield_chart"]

utilities/polars.py

@@ -714,7 +714,7 @@ class _DataClassToDataFrameNonUniqueError(DataClassToDataFrameError):
 
     @override
     def __str__(self) -> str:
-        return f"Iterable {get_repr(self.objs)} must contain exactly one class; got {self.first}, {self.second} and perhaps more"
+        return f"Iterable {get_repr(self.objs)} must contain exactly 1 class; got {self.first}, {self.second} and perhaps more"
 
 
 ##

utilities/polars_ols.py

@@ -1,19 +1,22 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, overload
 
-from polars import struct
+from polars import Expr, Series, struct
 from polars_ols import RollingKwargs, compute_rolling_least_squares
 
-from utilities.polars import ensure_expr_or_series
+from utilities.errors import ImpossibleCaseError
+from utilities.functions import is_sequence_of
+from utilities.polars import concat_series, ensure_expr_or_series
 
 if TYPE_CHECKING:
-    from polars import Expr
+    from polars._typing import IntoExprColumn
     from polars_ols import NullPolicy
 
     from utilities.polars import ExprLike
 
 
+@overload
 def compute_rolling_ols(
     target: ExprLike,
     *features: ExprLike,
@@ -24,9 +27,106 @@ def compute_rolling_ols(
     min_periods: int | None = None,
     use_woodbury: bool | None = None,
     alpha: float | None = None,
+) -> Expr: ...
+@overload
+def compute_rolling_ols(
+    target: Series,
+    *features: Series,
+    sample_weights: Series | None = None,
+    add_intercept: bool = False,
+    null_policy: NullPolicy = "drop_window",
+    window_size: int = 1000000,
+    min_periods: int | None = None,
+    use_woodbury: bool | None = None,
+    alpha: float | None = None,
+) -> Series: ...
+@overload
+def compute_rolling_ols(
+    target: IntoExprColumn,
+    *features: IntoExprColumn,
+    sample_weights: IntoExprColumn | None = None,
+    add_intercept: bool = False,
+    null_policy: NullPolicy = "drop_window",
+    window_size: int = 1000000,
+    min_periods: int | None = None,
+    use_woodbury: bool | None = None,
+    alpha: float | None = None,
+) -> Expr | Series: ...
+def compute_rolling_ols(
+    target: IntoExprColumn,
+    *features: IntoExprColumn,
+    sample_weights: IntoExprColumn | None = None,
+    add_intercept: bool = False,
+    null_policy: NullPolicy = "drop_window",
+    window_size: int = 1000000,
+    min_periods: int | None = None,
+    use_woodbury: bool | None = None,
+    alpha: float | None = None,
+) -> Expr | Series:
+    """Compute a rolling OLS."""
+    target = ensure_expr_or_series(target)
+    features2 = tuple(map(ensure_expr_or_series, features))
+    sample_weights = (
+        None if sample_weights is None else ensure_expr_or_series(sample_weights)
+    )
+    if (
+        isinstance(target, Expr)
+        and is_sequence_of(features2, Expr)
+        and ((sample_weights is None) or isinstance(sample_weights, Expr))
+    ):
+        return _compute_rolling_ols_expr(
+            target,
+            *features2,
+            sample_weights=sample_weights,
+            add_intercept=add_intercept,
+            null_policy=null_policy,
+            window_size=window_size,
+            min_periods=min_periods,
+            use_woodbury=use_woodbury,
+            alpha=alpha,
+        )
+    if (
+        isinstance(target, Series)
+        and is_sequence_of(features2, Series)
+        and ((sample_weights is None) or isinstance(sample_weights, Series))
+    ):
+        return concat_series(
+            target, *features2, *([] if sample_weights is None else [sample_weights])
+        ).with_columns(
+            _compute_rolling_ols_expr(
+                target.name,
+                *(f.name for f in features2),
+                sample_weights=None if sample_weights is None else sample_weights.name,
+                add_intercept=add_intercept,
+                null_policy=null_policy,
+                window_size=window_size,
+                min_periods=min_periods,
+                use_woodbury=use_woodbury,
+                alpha=alpha,
+            )
+        )["ols"]
+    raise ImpossibleCaseError(  # pragma: no cover
+        case=[f"{target=}", f"{features2=}", f"{sample_weights=}"]
+    )
+
+
+def _compute_rolling_ols_expr(
+    target: ExprLike,
+    *features: ExprLike,
+    sample_weights: ExprLike | None = None,
+    add_intercept: bool = False,
+    null_policy: NullPolicy = "drop_window",
+    window_size: int = 1000000,
+    min_periods: int | None = None,
+    use_woodbury: bool | None = None,
+    alpha: float | None = None,
 ) -> Expr:
     """Compute a rolling OLS."""
     target = ensure_expr_or_series(target)
+    features2 = tuple(map(ensure_expr_or_series, features))
+    sample_weights = (
+        None if sample_weights is None else ensure_expr_or_series(sample_weights)
+    )
     rolling_kwargs = RollingKwargs(
         null_policy=null_policy,
         window_size=window_size,
@@ -36,7 +136,7 @@ def compute_rolling_ols(
     )
     coefficients = compute_rolling_least_squares(
         target,
-        *features,
+        *features2,
         sample_weights=sample_weights,
         add_intercept=add_intercept,
         mode="coefficients",
@@ -44,7 +144,7 @@ def compute_rolling_ols(
     ).alias("coefficients")
     predictions = compute_rolling_least_squares(
         target,
-        *features,
+        *features2,
         sample_weights=sample_weights,
         add_intercept=add_intercept,
         mode="predictions",
@@ -52,7 +152,7 @@ def compute_rolling_ols(
     ).alias("predictions")
     residuals = compute_rolling_least_squares(
         target,
-        *features,
+        *features2,
         sample_weights=sample_weights,
         add_intercept=add_intercept,
         mode="residuals",