kabukit 0.1.1__py3-none-any.whl → 0.2.1__py3-none-any.whl

kabukit/__init__.py

@@ -1,4 +1,7 @@
+from .core.info import Info
+from .core.prices import Prices
+from .core.statements import Statements
 from .edinet.client import EdinetClient
 from .jquants.client import JQuantsClient
 
-__all__ = ["EdinetClient", "JQuantsClient"]
+__all__ = ["EdinetClient", "Info", "JQuantsClient", "Prices", "Statements"]

kabukit/analysis/visualization.py

@@ -0,0 +1,57 @@
+"""チャート作成のためのモジュール"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+import altair as alt
+
+if TYPE_CHECKING:
+    from kabukit.core.prices import Prices
+
+
+def plot_prices(
+    prices: Prices,
+    kind: Literal["candlestick"] = "candlestick",
+) -> alt.VConcatChart:
+    if kind == "candlestick":
+        chart_price = plot_prices_candlestick(prices)
+        chart_price_volume = plot_prices_volume(prices)
+        return alt.vconcat(chart_price, chart_price_volume)
+
+    raise NotImplementedError  # pyright: ignore[reportUnreachable]
+
+
+def plot_prices_candlestick(prices: Prices) -> alt.LayerChart:
+    rule = alt.Chart(prices.data, mark="rule").encode(y="Low:Q", y2="High:Q")
+    bar = alt.Chart(prices.data, mark="bar").encode(y="Open:Q", y2="Close:Q")
+
+    color_condition = alt.condition(
+        "datum.Open < datum.Close",
+        alt.value("#ff3030"),
+        alt.value("#3030ff"),
+    )
+
+    return alt.layer(rule, bar, height=200).encode(
+        x=alt.X("Date:T", axis=alt.Axis(title="日付", format="%Y-%m-%d")),
+        y=alt.Y(title="株価", scale=alt.Scale(zero=False)),
+        color=color_condition,
+        tooltip=[
+            alt.Tooltip("Date:T", title="日付"),
+            alt.Tooltip("Open:Q", title="始値"),
+            alt.Tooltip("High:Q", title="高値"),
+            alt.Tooltip("Low:Q", title="安値"),
+            alt.Tooltip("Close:Q", title="終値"),
+        ],
+    )
+
+
+def plot_prices_volume(prices: Prices) -> alt.Chart:
+    return alt.Chart(prices.data, mark="bar", height=50).encode(
+        x=alt.X("Date:T", axis=alt.Axis(title="日付", format="%Y-%m-%d")),
+        y=alt.Y("Volume:Q", title="出来高"),
+        tooltip=[
+            alt.Tooltip("Date:T", title="日付"),
+            alt.Tooltip("Volume:Q", title="出来高"),
+        ],
+    )

kabukit/cli/auth.py

@@ -50,7 +50,7 @@ async def jquants_alias(mailaddress: Mailaddress, password: Password) -> None:
 
 def auth_edinet(api_key: str) -> None:
     """EDINET APIのAPIキーを設定ファイルに保存します。"""
-    from kabukit.config import set_key
+    from kabukit.utils.config import set_key
 
     set_key("EDINET_API_KEY", api_key)
     typer.echo("EDINETのAPIキーを保存しました。")
@@ -75,7 +75,7 @@ def show() -> None:
     """設定ファイルに保存したトークン・APIキーを表示します。"""
     from dotenv import dotenv_values
 
-    from kabukit.config import get_dotenv_path
+    from kabukit.utils.config import get_dotenv_path
 
     path = get_dotenv_path()
     typer.echo(f"Configuration file: {path}")

kabukit/core/base.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import datetime
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import polars as pl
+from platformdirs import user_cache_dir
+
+if TYPE_CHECKING:
+    from typing import Self
+
+    from polars import DataFrame
+
+
+class Base:
+    data: DataFrame
+
+    def __init__(self, data: DataFrame) -> None:
+        self.data = data
+
+    @classmethod
+    def data_dir(cls) -> Path:
+        clsname = cls.__name__.lower()
+        return Path(user_cache_dir("kabukit")) / clsname
+
+    def write(self) -> Path:
+        data_dir = self.data_dir()
+        data_dir.mkdir(parents=True, exist_ok=True)
+        path = datetime.datetime.today().strftime("%Y%m%d")  # noqa: DTZ002
+        filename = data_dir / f"{path}.parquet"
+        self.data.write_parquet(filename)
+        return filename
+
+    @classmethod
+    def read(cls, path: str | None = None) -> Self:
+        data_dir = cls.data_dir()
+
+        if path is None:
+            filename = sorted(data_dir.glob("*.parquet"))[-1]
+        else:
+            filename = data_dir / path
+
+        data = pl.read_parquet(filename)
+        return cls(data)

kabukit/core/client.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from httpx import AsyncClient
+
+if TYPE_CHECKING:
+    from typing import Self
+
+
+class Client:
+    client: AsyncClient
+
+    def __init__(self, base_url: str = "") -> None:
+        self.client = AsyncClient(base_url=base_url)
+
+    async def aclose(self) -> None:
+        """HTTPクライアントを閉じる。"""
+        await self.client.aclose()
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:  # pyright: ignore[reportMissingParameterType, reportUnknownParameterType]  # noqa: ANN001
+        await self.aclose()

kabukit/core/info.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from .base import Base
+
+if TYPE_CHECKING:
+    from polars import DataFrame
+
+
+class Info(Base):
+    pass

kabukit/core/prices.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import polars as pl
+
+from .base import Base
+
+if TYPE_CHECKING:
+    from datetime import timedelta
+    from typing import Self
+
+    from polars import Expr
+
+
+class Prices(Base):
+    def truncate(self, every: str | timedelta | Expr) -> Self:
+        data = (
+            self.data.group_by(pl.col("Date").dt.truncate(every), "Code")
+            .agg(
+                pl.col("Open").drop_nulls().first(),
+                pl.col("High").max(),
+                pl.col("Low").min(),
+                pl.col("Close").drop_nulls().last(),
+                pl.col("Volume").sum(),
+                pl.col("TurnoverValue").sum(),
+            )
+            .sort("Code", "Date")
+        )
+        return self.__class__(data)

kabukit/core/statements.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from .base import Base
+
+
+class Statements(Base):
+    pass

kabukit/edinet/__init__.py

@@ -0,0 +1,3 @@
+from .concurrent import fetch, fetch_list
+
+__all__ = ["fetch", "fetch_list"]

kabukit/edinet/client.py

@@ -7,11 +7,13 @@ from enum import StrEnum
 from typing import TYPE_CHECKING
 
 import polars as pl
-from httpx import AsyncClient
 from polars import DataFrame
 
-from kabukit.config import load_dotenv
-from kabukit.params import get_params
+from kabukit.core.client import Client
+from kabukit.utils.config import load_dotenv
+from kabukit.utils.params import get_params
+
+from .doc import clean_csv, clean_list
 
 if TYPE_CHECKING:
     import datetime
@@ -29,11 +31,9 @@ class AuthKey(StrEnum):
     API_KEY = "EDINET_API_KEY"
 
 
-class EdinetClient:
-    client: AsyncClient
-
+class EdinetClient(Client):
     def __init__(self, api_key: str | None = None) -> None:
-        self.client = AsyncClient(base_url=BASE_URL)
+        super().__init__(BASE_URL)
         self.set_api_key(api_key)
 
     def set_api_key(self, api_key: str | None = None) -> None:
@@ -44,9 +44,6 @@ class EdinetClient:
         if api_key:
             self.client.params = {"Subscription-Key": api_key}
 
-    async def aclose(self) -> None:
-        await self.client.aclose()
-
     async def get(self, url: str, params: QueryParamTypes) -> Response:
         resp = await self.client.get(url, params=params)
         resp.raise_for_status()
@@ -71,7 +68,12 @@ class EdinetClient:
         if "results" not in data:
             return DataFrame()
 
-        return DataFrame(data["results"], infer_schema_length=None)
+        df = DataFrame(data["results"], infer_schema_length=None)
+
+        if df.is_empty():
+            return df
+
+        return clean_list(df, date)
 
     async def get_document(self, doc_id: str, doc_type: int) -> Response:
         params = get_params(type=doc_type)
@@ -105,7 +107,7 @@ class EdinetClient:
                             f.read(),
                             separator="\t",
                             encoding="utf-16-le",
-                        )
+                        ).pipe(clean_csv, doc_id)
 
         msg = "CSV is not available."
         raise ValueError(msg)

kabukit/edinet/concurrent.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+import datetime
+from typing import TYPE_CHECKING
+
+from kabukit.utils import concurrent
+
+from .client import EdinetClient
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from polars import DataFrame
+
+    from kabukit.utils.concurrent import Callback, Progress
+
+
+def get_dates(days: int | None = None, years: int | None = None) -> list[datetime.date]:
+    """過去days日またはyears年の日付リストを返す。
+
+    Args:
+        days (int | None): 過去days日の日付リストを取得する。
+        years (int | None): 過去years年の日付リストを取得する。
+            daysが指定されている場合は無視される。
+    """
+    end_date = datetime.date.today()  # noqa: DTZ011
+
+    if days is not None:
+        start_date = end_date - datetime.timedelta(days=days)
+    elif years is not None:
+        start_date = end_date.replace(year=end_date.year - years)
+    else:
+        msg = "daysまたはyearsのいずれかを指定してください。"
+        raise ValueError(msg)
+
+    return [
+        start_date + datetime.timedelta(days=i)
+        for i in range(1, (end_date - start_date).days + 1)
+    ]
+
+
+async def fetch(
+    resource: str,
+    args: Iterable[str | datetime.date],
+    /,
+    max_concurrency: int | None = None,
+    progress: Progress | None = None,
+    callback: Callback | None = None,
+) -> DataFrame:
+    """引数に対応する各種データを取得し、単一のDataFrameにまとめて返す。
+
+    Args:
+        resource (str): 取得するデータの種類。EdinetClientのメソッド名から"get_"を
+            除いたものを指定する。
+        args (Iterable[str | datetime.date]): 取得対象の引数のリスト。
+        max_concurrency (int | None, optional): 同時に実行するリクエストの最大数。
+            指定しないときはデフォルト値が使用される。
+        progress (Progress | None, optional): 進捗表示のための関数。
+            tqdm, marimoなどのライブラリを使用できる。
+            指定しないときは進捗表示は行われない。
+        callback (Callback | None, optional): 各DataFrameに対して適用する
+            コールバック関数。指定しないときはそのままのDataFrameが使用される。
+
+    Returns:
+        DataFrame:
+            すべての銘柄の財務情報を含む単一のDataFrame。
+    """
+    return await concurrent.fetch(
+        EdinetClient,
+        resource,
+        args,
+        max_concurrency=max_concurrency,
+        progress=progress,
+        callback=callback,
+    )
+
+
+async def fetch_list(
+    days: int | None = None,
+    years: int | None = None,
+    limit: int | None = None,
+    max_concurrency: int | None = None,
+    progress: Progress | None = None,
+    callback: Callback | None = None,
+) -> DataFrame:
+    """過去days日またはyears年の文書一覧を取得し、単一のDataFrameにまとめて返す。
+
+    Args:
+        days (int | None): 過去days日の日付リストを取得する。
+        years (int | None): 過去years年の日付リストを取得する。
+            daysが指定されている場合は無視される。
+        max_concurrency (int | None, optional): 同時に実行するリクエストの最大数。
+            指定しないときはデフォルト値が使用される。
+        progress (Progress | None, optional): 進捗表示のための関数。
+            tqdm, marimoなどのライブラリを使用できる。
+            指定しないときは進捗表示は行われない。
+        callback (Callback | None, optional): 各DataFrameに対して適用する
+            コールバック関数。指定しないときはそのままのDataFrameが使用される。
+
+    Returns:
+        DataFrame:
+            文書一覧を含む単一のDataFrame。
+    """
+    dates = get_dates(days=days, years=years)
+
+    if limit is not None:
+        dates = dates[:limit]
+
+    return await fetch(
+        "list",
+        dates,
+        max_concurrency=max_concurrency,
+        progress=progress,
+        callback=callback,
+    )
+
+
+async def fetch_csv(
+    doc_ids: Iterable[str],
+    /,
+    limit: int | None = None,
+    max_concurrency: int | None = None,
+    progress: Progress | None = None,
+    callback: Callback | None = None,
+) -> DataFrame:
+    """文書をCSV形式で取得し、単一のDataFrameにまとめて返す。
+
+    Args:
+        doc_ids (Iterable[str]): 取得対象の文書IDのリスト。
+        max_concurrency (int | None, optional): 同時に実行するリクエストの最大数。
+            指定しないときはデフォルト値が使用される。
+        progress (Progress | None, optional): 進捗表示のための関数。
+            tqdm, marimoなどのライブラリを使用できる。
+            指定しないときは進捗表示は行われない。
+        callback (Callback | None, optional): 各DataFrameに対して適用する
+            コールバック関数。指定しないときはそのままのDataFrameが使用される。
+
+    Returns:
+        DataFrame:
+            文書含む単一のDataFrame。
+    """
+    doc_ids = list(doc_ids)
+
+    if limit is not None:
+        doc_ids = doc_ids[:limit]
+
+    return await fetch(
+        "csv",
+        doc_ids,
+        max_concurrency=max_concurrency,
+        progress=progress,
+        callback=callback,
+    )

kabukit/edinet/doc.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import datetime
+from typing import TYPE_CHECKING
+
+import polars as pl
+
+if TYPE_CHECKING:
+    from polars import DataFrame
+
+
+def clean_list(df: DataFrame, date: str | datetime.date) -> DataFrame:
+    if isinstance(date, str):
+        date = datetime.datetime.strptime(date, "%Y-%m-%d").date()  # noqa: DTZ007
+
+    return df.with_columns(
+        pl.lit(date).alias("Date"),
+        pl.col("submitDateTime").str.to_datetime("%Y-%m-%d %H:%M", strict=False),
+        pl.col("^period.+$").str.to_date("%Y-%m-%d", strict=False),
+        pl.col("^.+Flag$").cast(pl.Int8).cast(pl.Boolean),
+        pl.col("^.+Code$").cast(pl.String),
+        pl.col("opeDateTime")
+        .cast(pl.String)
+        .str.to_datetime("%Y-%m-%d %H:%M", strict=False),
+    ).select("Date", pl.exclude("Date"))
+
+
+def clean_csv(df: DataFrame, doc_id: str) -> DataFrame:
+    return df.select(
+        pl.lit(doc_id).alias("docID"),
+        pl.all(),
+    )

kabukit/jquants/__init__.py

@@ -1,3 +1,3 @@
-from .stream import fetch, fetch_all
+from .concurrent import fetch, fetch_all
 
 __all__ = ["fetch", "fetch_all"]