Chore/docstrings and overrides#139
Conversation
Add Google-style docstrings to private functions, classes, and methods across compile.py, _ir.py, _ir_lower.py, _ir_templates.py, _axes.py, _vectorize.py, _errors.py, _normalize_ir.py, _normalize.py, _symbols.py, _templates.py, and _typing.py. No behavioral changes.
Add Examples sections with executable >>> blocks to compile_spec, compile_rhs, normalize_rhs, and normalize_expr_rhs.
Add the missing Returns: and Raises: sections required by ruff DOC2xx/ DOC5xx rules across the docstrings introduced earlier on this branch. Drop a few inaccurate Raises: claims on helpers that propagate via `_invalid()`. Also remove a stale FURB140/D102 noqa override by using itertools.starmap and adding a real EvalFn.__call__ docstring.
Extract small helpers from three functions so each falls under ruff's default cyclomatic-complexity threshold and the noqa: C901 directives can be removed: - _resolve_continuous_range_filter: extract _trapezoidal_weights_or_raise. - lower_subscript_to_buffer: extract _validate_wildcard_subscript_axes. - _build_aliases_ir_from_raw: extract _parse_alias_body_to_ir and _populate_alias_parsed_maps (with a small _AliasExpansionContext NamedTuple so the helper stays under PLR0913). No behavior change; full suite remains at 396 passing.
TimothyWillard
left a comment
TimothyWillard
left a comment
There was a problem hiding this comment.
I've seen this tuple[list[Any], int] (coords, size) type in a few places. Probably for a follow up but I think it could make sense to convert this to a named tuple/dataclass that is a generic over the list type. Especially since the size is a property of the coords.
Also I've noticed quite a few RST linking directives (e.g. :class:, :meth:, etc.). This project uses MkDocs for the documentation, not sphinx, so these will just be rendered as text. See http://127.0.0.1:8000/op_system/api-reference/compile/#op_system.compile.CompiledRhs.__setstate__ when viewing the documentation locally via just serve. There's mkdocs-autoref for linking, not sure if that's setup for this project or not.
Otherwise looks good, documentation improvements are always great. I think there's still room to add more examples to private helpers that will serve both as documentation and as doctests.
| @@ -64,6 +78,18 @@ def _normalize_axis_name(ax_map: Mapping[str, Any], *, idx: int, seen: set[str]) | |||
|
|
|||
|
|
|||
| def _normalize_axis_type(ax_map: Mapping[str, Any], *, idx: int) -> str: | |||
There was a problem hiding this comment.
| def _normalize_axis_type(ax_map: Mapping[str, Any], *, idx: int) -> str: | |
| def _normalize_axis_type(ax_map: Mapping[str, Any], *, idx: int) -> Literal["categorical", "continuous", "ordinal"]: |
Per the return type. May also make sense to convert this to a StrEnum in the future depending on use.
| @@ -705,6 +866,15 @@ def _expr_precedence(expr: Expr) -> int: | |||
|
|
|||
|
|
|||
| def _wrap(text: str, *, need: bool) -> str: | |||
There was a problem hiding this comment.
nit: usually when there's a boolean kwarg there's a correct default, I'm guessing maybe _wrap(text: str, *, need: bool = True) -> str based on the logic.
|
|
||
| Examples: | ||
| >>> import numpy as np | ||
| >>> compiled = compile_spec({ | ||
| ... "kind": "expr", | ||
| ... "state": ["x"], | ||
| ... "equations": {"x": "-x"}, | ||
| ... }) | ||
| >>> compiled.eval_fn(0.0, np.array([1.0])) | ||
| array([-1.]) |
There was a problem hiding this comment.
This is really nice, the doctest provides a clear example that is great for users/developers to quickly get an idea of what the function does either before/after reading the docstring itself.
| parsed_maps: Pair ``(reduce_parsed, full_parsed)`` of output maps | ||
| populated with reduce-form and point-expanded IRs. | ||
| """ | ||
| from op_system._ir_expand import expand_reduce_pointwise # noqa: PLC0415 |
There was a problem hiding this comment.
Is this to avoid a circular import? Usually I think of imports inside of a function as a code smell unless it's to guard against optional functionality which does not apply here.
| def _populate_alias_parsed_maps( | ||
| raw_name: str, | ||
| canonical_name: str, | ||
| ir_raw: Expr, | ||
| ctx: _AliasExpansionContext, | ||
| parsed_maps: tuple[dict[str, Expr], dict[str, Expr]], | ||
| ) -> None: |
There was a problem hiding this comment.
Does this perhaps make more sense as a method of _AliasExpansionContext rather than as a standalone function?
| Examples: | ||
| >>> import numpy as np | ||
| >>> from op_system.specs import normalize_rhs | ||
| >>> rhs = normalize_rhs({ | ||
| ... "kind": "expr", | ||
| ... "state": ["x"], | ||
| ... "equations": {"x": "2.0 * x"}, | ||
| ... }) | ||
| >>> compiled = compile_rhs(rhs) | ||
| >>> y = np.array([3.0]) | ||
| >>> compiled.eval_fn(0.0, y) | ||
| array([6.]) |
There was a problem hiding this comment.
Woo! Doctests are great!
2 participants
This pull request adds comprehensive docstrings to internal functions and classes throughout the
src/op_systempackage, significantly improving code readability and maintainability. The new docstrings clarify the purpose, arguments, return values, and exceptions for many utility functions, error classes, and IR (intermediate representation) handling routines. Additionally, an example usage was added to thecompile_specfunction.The most important changes are:
Documentation improvements for internal APIs:
src/op_system/_axes.py,src/op_system/_errors.py, andsrc/op_system/_ir.py, describing their purpose, arguments, return values, and exceptions. This includes validation, normalization, and parsing helpers for axis specifications and IR handling. [1] [2] [3] [4] [5] [6] [7] [8] [9] [10] [11] [12] [13] [14] [15] [16] [17] [18] [19] [20] [21] [22] [23] [24] [25]Example usage addition:
compile_specfunction docstring insrc/op_system/__init__.pyto illustrate its invocation and output.These improvements make the codebase more accessible to new contributors and facilitate easier debugging and extension.