diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 690ecdc..14d9cc0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,14 +1,6 @@ default_language_version: python: python3.10 repos: - - repo: https://github.com/pycqa/pydocstyle - rev: 6.3.0 - hooks: - - id: pydocstyle - files: ^beavers/(dag|replay|kafka|arrow).py - additional_dependencies: - - tomli - - repo: https://github.com/pre-commit/pre-commit-hooks rev: v5.0.0 hooks: diff --git a/beavers/__init__.py b/beavers/__init__.py index 1ac3cf4..00feaa9 100644 --- a/beavers/__init__.py +++ b/beavers/__init__.py @@ -1,3 +1,5 @@ +"""Main module for beavers package.""" + from beavers.dag import Dag, Node, TimerManager __version__ = "0.0.0" diff --git a/beavers/dag.py b/beavers/dag.py index 2828edb..5143d63 100644 --- a/beavers/dag.py +++ b/beavers/dag.py @@ -326,7 +326,8 @@ def _recalculate(self, cycle_id: int): except Exception as e: if self._frame_summaries: raise RuntimeError( - f"Unable to run node:\n{_format_frame_summaries(self._frame_summaries)}" + "Unable to run node:\n" + + _format_frame_summaries(self._frame_summaries) ) from e else: raise diff --git a/beavers/perspective_wrapper.py b/beavers/perspective_wrapper.py index 14fe55c..298f5cc 100644 --- a/beavers/perspective_wrapper.py +++ b/beavers/perspective_wrapper.py @@ -1,3 +1,5 @@ +"""DAG Wrapper to create a web application using perspective.""" + import dataclasses import pathlib from typing import Any, Literal, Optional, Sequence @@ -33,9 +35,7 @@ @dataclasses.dataclass(frozen=True) class PerspectiveTableDefinition: - """ - API table definition - """ + """API table definition.""" name: str index_column: str @@ -69,9 +69,7 @@ def validate(self, schema: pa.Schema): @dataclasses.dataclass(frozen=True) class _TableConfig: - """ - Internal perspective table config, which is passed to the html template - """ + """Internal perspective table config, which is passed to the html template.""" name: str index: str @@ -91,7 +89,7 @@ def from_definition(definition: PerspectiveTableDefinition, schema: pa.Schema): class TableRequestHandler(tornado.web.RequestHandler): - """Renders the table.html template, using the provided configurations""" + """Renders the table.html template, using the provided configurations.""" _tables: Optional[dict[str, _TableConfig]] = None _default_table: Optional[str] = None @@ -114,7 +112,7 @@ async def get(self, path: str) -> None: def _table_to_bytes(table: pa.Table) -> bytes: - """Serialize a table as bytes, to pass it to a perspective table""" + """Serialize a table as bytes, to pass it to a perspective table.""" with pa.BufferOutputStream() as sink: with pa.ipc.new_stream(sink, table.schema) as writer: for batch in table.to_batches(): @@ -137,7 +135,7 @@ class _PerspectiveNode: table: perspective.Table | None = None def __call__(self, table: pa.Table) -> None: - """Pass the arrow data to perspective""" + """Pass the arrow data to perspective.""" self.table.update(_table_to_bytes(table)) def get_table_config(self) -> _TableConfig: @@ -183,6 +181,7 @@ def to_perspective( def to_perspective_type(data_type: pa.DataType) -> Any: + """Convert a pyarrow DataType to a perspective type.""" for predicate, perspective_type in DATA_TYPES: if predicate(data_type): return perspective_type @@ -190,10 +189,11 @@ def to_perspective_type(data_type: pa.DataType) -> Any: def to_perspective_schema(schema: pa.Schema) -> dict[str, Any]: + """Convert a pyarrow Schema to a perspective schema.""" return {f.name: to_perspective_type(f.type) for f in schema} -def perspective_thread( +def _perspective_thread( perspective_server: perspective.Server, kafka_driver: KafkaDriver, nodes: list[_PerspectiveNode], @@ -218,6 +218,7 @@ def run_web_application( assets_directory: str = ASSETS_DIRECTORY, port: int = 8082, ) -> None: + """Run a tornado web application with perspective tables backed by a Beavers DAG.""" server = perspective.Server() nodes: list[_PerspectiveNode] = [] @@ -251,5 +252,5 @@ def run_web_application( ) web_app.listen(port) loop = tornado.ioloop.IOLoop.current() - loop.call_later(0, perspective_thread, server, kafka_driver, nodes) + loop.call_later(0, _perspective_thread, server, kafka_driver, nodes) loop.start() diff --git a/beavers/polars_wrapper.py b/beavers/polars_wrapper.py index 783766f..fdfc1dd 100644 --- a/beavers/polars_wrapper.py +++ b/beavers/polars_wrapper.py @@ -62,13 +62,12 @@ def source_table( self, schema: pl.Schema, name: Optional[str] = None ) -> Node[pl.DataFrame]: """Add a source stream of type `pl.DataFrame`.""" - return self._dag.source_stream(empty=schema.to_frame(), name=name) def table_stream( self, function: Callable[P, pl.DataFrame], schema: pl.Schema ) -> NodePrototype[pl.DataFrame]: - """Add a stream node of output type `pl.DataFrame`""" + """Add a stream node of output type `pl.DataFrame`.""" return self._dag.stream(function, empty=schema.to_frame()) def filter_stream( diff --git a/beavers/pyarrow_kafka.py b/beavers/pyarrow_kafka.py index dba31c0..235e765 100644 --- a/beavers/pyarrow_kafka.py +++ b/beavers/pyarrow_kafka.py @@ -1,3 +1,5 @@ +"""PyArrow Kafka serializers and deserializers.""" + import dataclasses import io import json @@ -15,6 +17,8 @@ @dataclasses.dataclass(frozen=True) class JsonDeserializer(KafkaMessageDeserializer[pa.Table]): + """Deserialize JSON messages from Kafka into a pyarrow Table.""" + schema: pa.Schema def __call__(self, messages: confluent_kafka.Message) -> pa.Table: @@ -36,6 +40,8 @@ def __call__(self, messages: confluent_kafka.Message) -> pa.Table: @dataclasses.dataclass(frozen=True) class JsonSerializer(KafkaMessageSerializer[pa.Table]): + """Serialize pyarrow Tables into JSON messages for Kafka.""" + topic: str def __call__(self, table: pa.Table): diff --git a/beavers/pyarrow_replay.py b/beavers/pyarrow_replay.py index 4f31cfa..66faea7 100644 --- a/beavers/pyarrow_replay.py +++ b/beavers/pyarrow_replay.py @@ -1,3 +1,5 @@ +"""Utils to replay historical data in a DAG.""" + import dataclasses from typing import Callable @@ -9,6 +11,8 @@ class ArrowTableDataSource(DataSource[pa.Table]): + """A replay data source that replay data from an arrow Table.""" + def __init__( self, table: pa.Table, timestamp_extractor: Callable[[pa.Table], pa.Array] ): @@ -42,6 +46,8 @@ def get_next(self) -> pd.Timestamp: @dataclasses.dataclass class ArrowTableDataSink(DataSink[pa.Table]): + """A data sink that save data in an arrow Table.""" + saver: Callable[[pa.Table], None] chunks: list[pa.Table] = dataclasses.field(default_factory=list) diff --git a/beavers/pyarrow_wrapper.py b/beavers/pyarrow_wrapper.py index 4605380..c4dc038 100644 --- a/beavers/pyarrow_wrapper.py +++ b/beavers/pyarrow_wrapper.py @@ -105,7 +105,7 @@ def source_table( def table_stream( self, function: Callable[P, pa.Table], schema: pa.Schema ) -> NodePrototype[pa.Table]: - """Add a stream node of output type `pa.Table`""" + """Add a stream node of output type `pa.Table`.""" return self._dag.stream(function, empty=schema.empty_table()) def filter_stream( diff --git a/beavers/replay.py b/beavers/replay.py index 3f30f6d..54455b7 100644 --- a/beavers/replay.py +++ b/beavers/replay.py @@ -21,7 +21,7 @@ class ReplayContext: """ Stores the information about a replay. - Attributes + Attributes: ---------- start: pd.Timestamp Start of the replay @@ -58,7 +58,7 @@ def read_to(self, timestamp: pd.Timestamp) -> T: timestamp End of the time interval for which data is required (inclusive) - Returns + Returns: ------- data The data for the interval (or empty if no data is found) @@ -72,7 +72,7 @@ def get_next(self) -> pd.Timestamp: If no data is available this should return `UTC_MAX` - Returns + Returns: ------- timestamp: pd.Timestamp Timestamp of the next available data point (or `UTC_MAX` if no more data @@ -113,7 +113,7 @@ def __call__(self, replay_context: ReplayContext) -> DataSource[T]: replay_context: Information about the replay that's about to run - Returns + Returns: ------- DataSource[T]: Source for the replay @@ -134,7 +134,7 @@ def __call__(self, replay_context: ReplayContext) -> DataSink[T]: replay_context: Information about the replay that's about to run - Returns + Returns: ------- DataSink[T]: Sink for the replay @@ -187,7 +187,7 @@ class ReplayDriver: - collect the output data and pass it to the sink - close the sink at the end of the run - Notes + Notes: ----- Do not call the constructor directly, use `create` instead diff --git a/beavers/testing.py b/beavers/testing.py index 2cf5d1f..62a19f1 100644 --- a/beavers/testing.py +++ b/beavers/testing.py @@ -1,3 +1,5 @@ +"""Helper functions and tool to test a beavers DAG.""" + from typing import Any, Optional, Sequence, TypeVar import pandas as pd @@ -8,6 +10,8 @@ class DagTestBench: + """A test bench to test a beavers DAG.""" + def __init__(self, dag: Dag): self.dag = dag for output_name, output_sinks in self.dag.get_sinks().items(): diff --git a/pyproject.toml b/pyproject.toml index 1c14f75..52daaaf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -110,5 +110,17 @@ asyncio_mode = "auto" [tool.ruff] line-length = 88 +[tool.ruff.lint] +ignore = ["D102", "D107", "D203", "D212"] +select = ["D", "E", "F"] + [tool.ruff.lint.isort] known-first-party = ["beavers", "tradewell_proto"] + +[tool.ruff.lint.per-file-ignores] +"examples/*" = ["D"] +"scripts/*" = ["D"] +"tests/*" = ["D"] + +[tool.ruff.lint.pydocstyle] +convention = "google" diff --git a/tests/test_perpective_wrapper.py b/tests/test_perpective_wrapper.py index 2b334db..6f430bb 100644 --- a/tests/test_perpective_wrapper.py +++ b/tests/test_perpective_wrapper.py @@ -17,7 +17,7 @@ _table_to_bytes, _TableConfig, _UpdateRunner, - perspective_thread, + _perspective_thread, ) PERSPECTIVE_TABLE_SCHEMA = pa.schema( @@ -118,7 +118,7 @@ def start(self): def test_perspective_thread(): manager = Server() - perspective_thread(manager, MagicMock(), []) + _perspective_thread(manager, MagicMock(), []) class TestHandler(AsyncHTTPTestCase):