Add ruff for pydoc style

.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:

beavers/__init__.py

@@ -1,3 +1,5 @@
+"""Main module for beavers package."""
+
 from beavers.dag import Dag, Node, TimerManager
 
 __version__ = "0.0.0"

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

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,17 +181,19 @@ 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
     raise TypeError(f"Unsupported type: {data_type}")
 
 
 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()

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(

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):

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)
 

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(

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
 

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():

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"

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):