Add mypy

README.md

@@ -8,12 +8,17 @@
 [![License][license-image]][license-url]
 [![Downloads][downloads-image]][downloads-url]
 [![Downloads][downloads-month-image]][downloads-month-url]
-[![Code style: black][codestyle-image]][codestyle-url]
 [![snyk][snyk-image]][snyk-url]
+[![pydocstyle][pydocstyle-image]][pydocstyle-url]
+[![Checked with mypy][mypy-image]][mypy-url]
+[![Ruff][ruff-image]][ruff-url]
+[![pre-commit enable][precommit-image]][precommit-url]
+[![semantic-release][semver-image]][semver-url]
 <a href="https://trackgit.com">
 <img src="https://us-central1-trackgit-analytics.cloudfunctions.net/token/ping/m7c1fpbnueo78pkcetcm" alt="trackgit-views" />
 </a>
 
+
 ![Beavers Logo][5]
 
 # Beavers
@@ -80,7 +85,15 @@ for both realtime and batch jobs.
 [downloads-url]: https://static.pepy.tech/badge/beavers
 [downloads-month-image]: https://pepy.tech/badge/beavers/month
 [downloads-month-url]: https://static.pepy.tech/badge/beavers/month
-[codestyle-image]: https://img.shields.io/badge/code%20style-black-000000.svg
-[codestyle-url]: https://github.com/ambv/black
 [snyk-image]: https://snyk.io/advisor/python/beavers/badge.svg
 [snyk-url]: https://snyk.io/advisor/python/beavers
+[pydocstyle-image]: https://img.shields.io/badge/pydocstyle-enabled-AD4CD3
+[pydocstyle-url]: http://www.pydocstyle.org/en/stable/
+[mypy-image]: http://www.mypy-lang.org/static/mypy_badge.svg
+[mypy-url]: http://mypy-lang.org/
+[ruff-image]: https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json
+[ruff-url]: https://github.com/astral-sh/ruff
+[precommit-image]: https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white
+[precommit-url]: https://pre-commit.com/
+[semver-image]: https://img.shields.io/badge/%20%20%F0%9F%93%A6%F0%9F%9A%80-semantic--release-e10079.svg
+[semver-url]: https://semver.org/

beavers/dag.py

@@ -1,6 +1,7 @@
 """Module for building and representing dags and nodes."""
 
 from __future__ import annotations
+import pandas as pd
 
 import collections.abc
 import dataclasses
@@ -23,21 +24,20 @@
     try:
         from beavers.pyarrow_wrapper import ArrowDagWrapper
     except ImportError:
-        ArrowDagWrapper = None
+        ArrowDagWrapper = None  # type: ignore[assignment, misc]
     try:
         from beavers.pandas_wrapper import PandasWrapper
     except ImportError:
-        PandasWrapper = None
+        PandasWrapper = None  # type: ignore[assignment, misc]
     try:
         from beavers.perspective_wrapper import PerspectiveDagWrapper
     except ImportError:
-        PerspectiveDagWrapper = None
+        PerspectiveDagWrapper = None  # type: ignore[assignment, misc]
     try:
         from beavers.polars_wrapper import PolarsDagWrapper
     except ImportError:
-        PolarsDagWrapper = None
+        PolarsDagWrapper = None  # type: ignore[assignment, misc]
 
-import pandas as pd
 
 P = ParamSpec("P")
 
@@ -58,7 +58,7 @@ def _format_frame_summaries(frame_summaries: Sequence[traceback.FrameSummary]) -
 
 
 class _SourceStreamFunction(Generic[T]):
-    def __init__(self, empty_factory: Callable[[], T], name: str):
+    def __init__(self, empty_factory: Callable[[], T], name: str | None):
         self._empty_factory = empty_factory
         self._name = name
         self._value = empty_factory()
@@ -87,15 +87,15 @@ def __call__(self, value: T) -> None:
 
 class _ValueCutOff(Generic[T]):
     def __init__(self, comparator: Callable[[T, T], bool] = operator.eq):
-        self._value = _STATE_EMPTY
+        self._value: T = _STATE_EMPTY  # type: ignore[assignment]
         self._comparator = comparator
 
     def __call__(self, value: T) -> T:
-        if self._value == _STATE_EMPTY or not self._comparator(value, self._value):
+        if self._value is _STATE_EMPTY or not self._comparator(value, self._value):
             self._value = value
             return value
         else:
-            return _STATE_UNCHANGED
+            return _STATE_UNCHANGED  # type: ignore[return-value]
 
 
 class TimerManager:
@@ -111,7 +111,7 @@ class TimerManager:
     It is accessed by the framework to decide when the next timer.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize with default values."""
         self._next_timer: pd.Timestamp = UTC_MAX
         self._just_triggered: bool = False
@@ -153,7 +153,7 @@ class _TimerManagerFunction:
     Used by the framework to trigger timers.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.timer_manager: TimerManager = TimerManager()
 
     def __call__(self) -> TimerManager:
@@ -210,7 +210,7 @@ def create(positional: Sequence[Node], key_word: dict[str, Node]) -> "_NodeInput
             nodes=tuple(all_nodes),
         )
 
-    def input_nodes(self) -> tuple[Node]:
+    def input_nodes(self) -> tuple[Node, ...]:
         return self.nodes
 
 
@@ -242,17 +242,23 @@ class Node(Generic[T]):
     You shouldn't use them directly to read values (use sink for this)
     """
 
-    _function: Optional[Callable[..., T]]
+    _function: Callable[..., T]
     _inputs: _NodeInputs = dataclasses.field(repr=False)
     _empty_factory: Any
     _observers: list[Node] = dataclasses.field(repr=False)
     _runtime_data: _RuntimeNodeData
     _frame_summaries: tuple[traceback.FrameSummary, ...] = ()
 
+    def __post_init__(self):
+        """Check not is valid."""
+        assert self._function is not None
+        assert callable(self._function)
+
     @staticmethod
     def _create(
-        value: T = None,
-        function: Optional[Callable[..., T]] = None,
+        value: T | None = None,
+        *,
+        function: Callable[..., T],
         inputs: _NodeInputs = _NO_INPUTS,
         empty_factory: Any = _STATE_EMPTY,
         notifications: int = 1,
@@ -272,7 +278,7 @@ def get_value(self) -> T:
         if self._runtime_data.value is _VALUE_EMPTY:
             return self._empty_factory()
         else:
-            return self._runtime_data.value
+            return self._runtime_data.value  # type: ignore[return-value]
 
     def get_cycle_id(self) -> int:
         """Return id of the cycle at which this node last updated."""
@@ -282,14 +288,16 @@ def set_stream(self, value: T):
         """Set the value of a `_SourceStream`."""
         if not isinstance(self._function, _SourceStreamFunction):
             raise TypeError(f"Only {_SourceStreamFunction.__name__} can be set")
-        self._function.set(value)
-        self._stain()
+        else:
+            self._function.set(value)
+            self._stain()
 
     def get_sink_value(self) -> Any:
         """Return the value of a `_SinkFunction`."""
-        if not isinstance(self._function, _SinkFunction):
+        if isinstance(self._function, _SinkFunction):
+            return self._function.get()  # type: ignore[attr-defined]
+        else:
             raise TypeError(f"Only {_SinkFunction.__name__} can be read")
-        return self._function.get()
 
     def _stain(self):
         self._runtime_data.notifications += 1
@@ -397,13 +405,13 @@ class DagMetrics:
 class Dag:
     """Main class used for building and executing a dag."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Create an empty `Dag`."""
         self._nodes: list[Node] = []
         self._sources: dict[str, Node] = {}
         self._sinks: dict[str, Node] = {}
         self._now_node: Node[pd.Timestamp] = self._add_node(
-            Node._create(UTC_EPOCH, _SourceState(UTC_EPOCH))
+            Node._create(value=UTC_EPOCH, function=_SourceState(UTC_EPOCH))
         )
         self._silent_now_node: Node[pd.Timestamp] = self.silence(self._now_node)
         self._timer_manager_nodes: list[Node[TimerManager]] = []
@@ -517,7 +525,7 @@ def state(self, function: Callable[P, T]) -> NodePrototype[T]:
         _check_function(function)
 
         def add_to_dag(
-            inputs: _NodeInputs, frame_summaries: tuple[traceback.FrameSummary]
+            inputs: _NodeInputs, frame_summaries: list[traceback.FrameSummary]
         ) -> Node:
             return self._add_state(function, inputs, frame_summaries)
 
@@ -604,7 +612,7 @@ def silence(self, node: Node[T]) -> Node[T]:
         _check_input(node)
         return self._add_node(
             Node._create(
-                function=SilentUpdate,
+                function=SilentUpdate,  # type: ignore[arg-type]
                 inputs=_NodeInputs.create([node], {}),
                 value=node.get_value(),
                 empty_factory=node._empty_factory,
@@ -670,7 +678,7 @@ def execute(self, timestamp: Optional[pd.Timestamp] = None):
         """Run the dag for a given timestamp."""
         self._cycle_id += 1
         if timestamp is not None:
-            self._now_node._function.set_value(timestamp)
+            self._now_node._function.set_value(timestamp)  # type: ignore[attr-defined]
             self._now_node._stain()
             self._flush_timers(timestamp)
 
@@ -736,10 +744,13 @@ def _add_stream(
             )
         )
 
-    def _flush_timers(self, now: pd.Timestamp) -> int:
+    def _flush_timers(
+        self,
+        now: pd.Timestamp,  # type: ignore[name-defined]
+    ) -> int:
         count = 0
         for node in self._timer_manager_nodes:
-            timer_manager: TimerManager = node.get_value()
+            timer_manager: TimerManager = node.get_value()  # type: ignore[name-defined]
             if timer_manager._flush(now):
                 node._stain()
                 count += 1
@@ -778,14 +789,14 @@ def _check_empty(
     if empty is not None and empty_factory is not None:
         raise ValueError(f"Can't provide both {empty=} and {empty_factory=}")
     elif empty is None and empty_factory is None:
-        return list
+        return list  # type: ignore[return-value]
     elif empty is not None:
         if not isinstance(empty, collections.abc.Sized):
             raise TypeError("`empty` should implement `__len__`")
         elif len(empty) != 0:
             raise TypeError("`len(empty)` should be 0")
         else:
-            return lambda: empty
+            return lambda: empty  # type: ignore[return-value]
     else:
         assert empty is None
         if not callable(empty_factory):
@@ -809,7 +820,7 @@ def _check_input(node: Node) -> Node:
         return node
 
 
-def _check_function(function: Callable[ParamSpec, T]) -> Callable[ParamSpec, T]:
+def _check_function(function: Callable[P, T]) -> Callable[P, T]:
     if not callable(function):
         raise TypeError("`function` should be a callable")
     else:

beavers/kafka.py

@@ -247,7 +247,7 @@ class _ConsumerManager:
     def __init__(
         self,
         cutoff: pd.Timestamp,
-        partitions: dict[confluent_kafka.TopicPartition : tuple[int, int]],
+        partitions: dict[confluent_kafka.TopicPartition, tuple[int, int]],
         consumer: confluent_kafka.Consumer,
         batch_size: int,
         max_held_messages: int,

docs/contributing.md

@@ -4,6 +4,8 @@ Welcome! We're happy to have you here. Thank you in advance for your contributio
 
 ## Development environment set up
 
+The repo uses poetry to manage dependencies.
+
 ```shell
 python3 -m venv --clear venv
 source venv/bin/activate
@@ -27,12 +29,38 @@ coverage run --branch --rcfile=./pyproject.toml --include "./beavers/*" -m pytes
 coverage report --show-missing
 ```
 
+To run on every python version with Tox:
+
+```shell
+tox run
+```
+
+## Linting
+
+The repo uses [pre-commit](https://pre-commit.com/) to lint automatically.
+To install pre-commit hooks, which will lint automatically on commits:
+```shell
+pre-commit install # Run once to run automatically when committing
+```
+
+To run the linter:
+
+```shell
+pre-commit run --all-files
+```
+We started using mypy, on a limited number of files only for the moment.
+Mypy doesn't work in pre-commit, so you need to run it manually:
+
+```shell
+mypy --config-file=./pyproject.toml ./beavers/dag.py
+```
+
+
 ## Generating the change log
 
 We use [git-change-log](https://pawamoy.github.io/git-changelog/usage/) to generate our CHANGELOG.md
 
-Please follow the [basic convention](https://pawamoy.github.io/git-changelog/usage/#basic-convention) for commit
-message.
+Please follow the [basic convention](https://pawamoy.github.io/git-changelog/usage/#basic-convention) for commit message.
 
 To update the change log, run:
 
@@ -55,7 +83,7 @@ git tag vX.X.X
 git push origin vX.X.X
 ```
 
-Lastly on github, go to tags and create a release.
+Lastly on [Github](https://github.com/tradewelltech/beavers), go to tags and create a release.
 The CI will deploy to pypi automatically from then.
 
 ## Testing the documentation