Skip to content

Ensure runtime type-checking with beartype #143

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 97 commits into
base: main
Choose a base branch
from
Open

Ensure runtime type-checking with beartype #143

wants to merge 97 commits into from

Conversation

adamamer20
Copy link
Member

@adamamer20 adamamer20 commented Apr 4, 2025

This PR adds beartype to enable runtime type checking across the codebase.

Given the extensive use of type hints already present, the integration should be straightforward and introduce minimal performance overhead. Preliminary testing on the sugarscape example showed negligible impact, though broader testing across the full codebase hasn't been performed yet.

Runtime type checking provides several benefits:

  1. Simplifies debugging and improves development experience for both contributors and users.
  2. Ensures consistency between type annotations and actual runtime behavior.
  3. Reduces the need to manually handle type-related errors in our own logic.

Some notes regarding this PR:

  • I removed typing_extensions and changed to typing with the support of Python bumped up to 3.11. This is in line with the current mesa repository.
  • Following the advice of beartype, I completely removed if TYPE_CHECKING blocks for forward imports. This in fact doesn't work well with beartype. Instead, using from __future__ import annotations at the start of the file allows for the postponed evaluation of types
  • Forward imports do not use a string (also because it is upgraded by pyupgrade). This may result in a warning by Ruff but it's not an issue.

Performance considerations

I ran a side-by-side comparison of performance_plot.py with and without beartype (on Codespaces) to quantify the overhead. Below is the full picture:

Key points (0 – 100 k agents)

Agents pl concise pl native
0 +133 % +133 %
10 000 +21 % +67 %
50 000 0 % +35 %
90 000 +3 % +6 %
100 000 +113 % +261 %

Larger runs (100 k – 1 000 k agents)

Agents pl concise pl native
200 000 +2 % −7 % (7 % faster)
500 000 +30 % +8 %
1 000 000 +13 % +5 %
  • Startup overhead at 0 agents is ~133 %, mostly due to module import and GC initialization.

  • For small→mid runs (10 k–90 k), expect +20–40 % on the pandas-concise path and +5–70 % on native-Polars.

  • A sharp spike at exactly 100 k agents (+113 % / +261 %) arises from crossing Python’s GC and list-resizing thresholds under heavy temporary-object churn.

  • Beyond 200 k, the relative penalty subsides to +2–30 % on concise and −7 %→+8 % on native.

  • Disable runtime checks with -O.
    When you run Python under the -O (optimize) flag, all @beartype checks are skipped—so you recover essentially the same performance as the “no-beartype” build.

    python -O performance_plot.py
    

    (See the discussion in the Beartype issue tracker for details:
    [FAQ] Document how to disable @beartype beartype/beartype#94)

  • Summary by CodeRabbit

    • Refactor

      • Enhanced type annotations and added beartype runtime type checking to core classes for improved type safety.
      • Expanded and refined type aliases to support broader input types and optional dependencies.
      • Updated method signatures across modules for greater flexibility and consistency.
      • Improved documentation clarity and standardized import statements.
    • Chores

      • Replaced typeguard with beartype for runtime type checking in tests and dependencies.
      • Increased minimum Python version requirement to 3.11.
      • Updated .gitignore and pre-commit configurations.
    • Tests

      • Updated test decorators to use beartype for runtime type validation.
      • Minor adjustments in test data construction for compatibility.

Copy link
Contributor

coderabbitai bot commented Apr 4, 2025

Important

Review skipped

Auto reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Walkthrough

This update introduces comprehensive enhancements to type annotations, runtime type checking, and type alias definitions across the codebase. The beartype library is now used for runtime type validation, replacing typeguard in both source and test code. Type hints are made more precise and expressive, with many method signatures updated to accept broader or more accurate input types. Numerous classes are now decorated with @beartype for runtime enforcement. Type aliases are expanded to support additional input types and stricter constraints. Minor import reorganizations and docstring clarifications are included, but no core logic or control flow is altered.

Changes

File(s) Change Summary
.gitignore Added llm_rules.md to the ignore list.
pyproject.toml Removed typing-extensions dependency; replaced typeguard with beartype in test dependencies; added pre-commit to dev dependencies.
mesa_frames/types_.py Added from __future__ import annotations; expanded imports; added/extended type aliases for Polars, DataFrame, Index, Series, Infinity, etc.; improved compatibility with optional dependencies.
mesa_frames/abstract/agents.py Refined type annotations to use fully qualified concrete types; added @beartype to AgentContainer and AgentSetDF; extended method signatures for type precision and runtime validation; improved indexing capabilities.
mesa_frames/abstract/mixin.py Added from __future__ import annotations; added @beartype to CopyMixin and DataFrameMixin; refined type annotations for method parameters; improved typing consistency.
mesa_frames/abstract/space.py Added from __future__ import annotations; added @beartype to SpaceDF, DiscreteSpaceDF, GridDF; improved type hints, especially for model references and method parameters; expanded accepted input types.
mesa_frames/concrete/agents.py Added from __future__ import annotations; added @beartype to AgentsDF; updated type annotations to use concrete types; improved docstrings and method signatures for clarity and precision.
mesa_frames/concrete/agentset.py Added from __future__ import annotations; reorganized imports; added @beartype to AgentSetPolars; updated constructor type annotation for model.
mesa_frames/concrete/mixin.py Added from __future__ import annotations; added @beartype to PolarsMixin; generalized method parameter types from Sequence to Collection; improved overloads and type signatures for DataFrame/Series operations.
mesa_frames/concrete/model.py Added from __future__ import annotations; added @beartype to ModelDF; updated method return types to use concrete class references; documented setter parameter.
mesa_frames/concrete/space.py Added @beartype to GridPolars; changed capacity handling to accept None for unlimited; updated types for methods/properties related to capacity and agent input.
mesa_frames/__init__.py Added from __future__ import annotations; reordered imports for clarity.
tests/test_agentset.py, tests/test_grid.py, tests/test_mixin.py Replaced typeguard decorators with beartype on test classes.
tests/test_agents.py Changed creation of Polars "unique_id" column from pl.lit() to pl.Series() in test.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant ModelDF
    participant AgentSetDF
    participant SpaceDF

    User->>ModelDF: create ModelDF()
    ModelDF->>AgentSetDF: create AgentSetDF(model)
    AgentSetDF->>SpaceDF: interact with SpaceDF (typed as concrete)
    AgentSetDF->>AgentSetDF: add/discard/remove agents (runtime checked by @beartype)
    AgentSetDF->>AgentSetDF: get/set agents, use advanced indexing
    ModelDF->>AgentSetDF: get_agents_of_type(type) returns AgentSetDF
Loading

Poem

🐇
With beartype's hug, our types are snug,
From abstract roots to concrete rug.
Hints now precise, and checks at run,
Imports aligned, the typing's fun!
The rabbit hops through fields anew,
Polars and pandas, all in view.
Code is tidy, strong, and bright—
Happy typing, day and night!


Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share
🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Generate unit testing code for this file.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
    • @coderabbitai generate unit testing code for this file.
    • @coderabbitai modularize this function.
  • PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read src/utils.ts and generate unit testing code.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
    • @coderabbitai help me debug CodeRabbit configuration file.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

  • @coderabbitai pause to pause the reviews on a PR.
  • @coderabbitai resume to resume the paused reviews.
  • @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
  • @coderabbitai full review to do a full review from scratch and review all the files again.
  • @coderabbitai summary to regenerate the summary of the PR.
  • @coderabbitai generate docstrings to generate docstrings for this PR.
  • @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
  • @coderabbitai resolve resolve all the CodeRabbit review comments.
  • @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
  • @coderabbitai help to get help.

Other keywords and placeholders

  • Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

  • You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
  • Please see the configuration documentation for more information.
  • If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

@adamamer20 adamamer20 marked this pull request as draft April 4, 2025 07:25
@adamamer20 adamamer20 self-assigned this Apr 4, 2025
@adamamer20 adamamer20 added the feature New functionality added to the project. label Apr 4, 2025
@adamamer20 adamamer20 added this to the 0.1.0-stable milestone Apr 4, 2025
@adamamer20 adamamer20 linked an issue Apr 4, 2025 that may be closed by this pull request
@adamamer20 adamamer20 marked this pull request as ready for review April 29, 2025 08:22
@adamamer20
Copy link
Member Author

@coderabbitai review

Copy link
Contributor

coderabbitai bot commented Apr 29, 2025

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 5

🔭 Outside diff range comments (1)
mesa_frames/abstract/agents.py (1)

48-55: ⚠️ Potential issue

Import guard for Self (build blocker)

Same pattern required here:

-from typing import Self, overload, Any, Literal
+from typing import overload, Any, Literal, TYPE_CHECKING
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+if TYPE_CHECKING:
+    from mesa_frames.concrete.model import ModelDF
+    from mesa_frames.concrete.agents import AgentSetDF

This removes the current CI failure and also lets you shorten many very long
annotations (mesa_frames.concrete.agents.AgentSetDF) to the alias
AgentSetDF, improving readability.

🧹 Nitpick comments (2)
mesa_frames/abstract/space.py (2)

100-108: Constructor annotation may trigger circular import

beartype will import mesa_frames.concrete.model at runtime. To stay safe:

  1. Add the TYPE_CHECKING guarded import (see above).
  2. Replace the long annotation with the local alias ModelDF.

No behaviour change, fewer import cycles.

🧰 Tools
🪛 Ruff (0.8.2)

100-100: Undefined name mesa_frames

(F821)


1108-1116: Docstring no longer matches type

remaining_capacity now returns int | Infinity but docstring still says
“None if the capacity is infinite”. Update wording to reflect the new sentinel
value (math.inf / Infinity) so users aren’t confused.

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 5078421 and de7da65.

⛔ Files ignored due to path filters (1)
  • uv.lock is excluded by !**/*.lock
📒 Files selected for processing (16)
  • .gitignore (1 hunks)
  • mesa_frames/__init__.py (1 hunks)
  • mesa_frames/abstract/agents.py (27 hunks)
  • mesa_frames/abstract/mixin.py (8 hunks)
  • mesa_frames/abstract/space.py (17 hunks)
  • mesa_frames/concrete/agents.py (12 hunks)
  • mesa_frames/concrete/agentset.py (2 hunks)
  • mesa_frames/concrete/mixin.py (15 hunks)
  • mesa_frames/concrete/model.py (3 hunks)
  • mesa_frames/concrete/space.py (4 hunks)
  • mesa_frames/types_.py (3 hunks)
  • pyproject.toml (1 hunks)
  • tests/test_agents.py (1 hunks)
  • tests/test_agentset.py (1 hunks)
  • tests/test_grid.py (2 hunks)
  • tests/test_mixin.py (1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (3)
tests/test_mixin.py (1)
mesa_frames/concrete/mixin.py (1)
  • PolarsMixin (58-649)
mesa_frames/concrete/agentset.py (5)
mesa_frames/concrete/model.py (3)
  • agents (163-181)
  • agents (184-187)
  • ModelDF (56-228)
mesa_frames/abstract/agents.py (6)
  • agents (698-704)
  • agents (708-716)
  • agents (1102-1103)
  • agents (1106-1114)
  • AgentSetDF (783-1135)
  • model (667-674)
mesa_frames/concrete/agents.py (2)
  • agents (558-559)
  • agents (562-570)
mesa_frames/concrete/mixin.py (1)
  • PolarsMixin (58-649)
mesa_frames/utils.py (1)
  • copydoc (4-18)
mesa_frames/abstract/agents.py (7)
mesa_frames/abstract/space.py (4)
  • random (532-539)
  • model (522-529)
  • agents (512-519)
  • SpaceDF (89-539)
mesa_frames/abstract/mixin.py (2)
  • CopyMixin (66-184)
  • _df_get_masked_df (328-335)
tests/test_grid.py (1)
  • model (22-29)
mesa_frames/concrete/model.py (5)
  • ModelDF (56-228)
  • agents (163-181)
  • agents (184-187)
  • space (201-218)
  • space (221-228)
mesa_frames/concrete/agents.py (12)
  • agents (558-559)
  • agents (562-570)
  • do (147-155)
  • do (158-166)
  • do (168-203)
  • inactive_agents (605-606)
  • index (609-610)
  • pos (613-614)
  • add (87-120)
  • remove (216-260)
  • active_agents (573-574)
  • active_agents (577-580)
mesa_frames/concrete/agentset.py (8)
  • agents (498-499)
  • agents (502-505)
  • inactive_agents (516-517)
  • index (520-521)
  • pos (524-525)
  • add (100-149)
  • active_agents (508-509)
  • active_agents (512-513)
mesa_frames/concrete/mixin.py (1)
  • _df_get_masked_df (325-336)
🪛 Ruff (0.8.2)
mesa_frames/concrete/agentset.py

67-67: typing.Literal imported but unused

Remove unused import: typing.Literal

(F401)


71-71: mesa_frames.concrete.model.ModelDF imported but unused

Remove unused import: mesa_frames.concrete.model.ModelDF

(F401)


72-72: mesa_frames.types_.IdsLike imported but unused

Remove unused import: mesa_frames.types_.IdsLike

(F401)


88-88: Undefined name mesa_frames

(F821)

mesa_frames/concrete/agents.py

75-75: Undefined name mesa_frames

(F821)

mesa_frames/abstract/space.py

100-100: Undefined name mesa_frames

(F821)


522-522: Undefined name mesa_frames

(F821)


555-555: Undefined name mesa_frames

(F821)


1157-1157: Undefined name mesa_frames

(F821)

mesa_frames/abstract/agents.py

71-71: Undefined name mesa_frames

(F821)


80-80: Undefined name mesa_frames

(F821)


81-81: Undefined name mesa_frames

(F821)


107-107: Undefined name mesa_frames

(F821)


108-108: Undefined name mesa_frames

(F821)


134-134: Undefined name mesa_frames

(F821)


139-139: Undefined name mesa_frames

(F821)


176-176: Undefined name mesa_frames

(F821)


187-187: Undefined name mesa_frames

(F821)


250-250: Undefined name mesa_frames

(F821)


251-251: Undefined name mesa_frames

(F821)


399-399: Undefined name mesa_frames

(F821)


400-400: Undefined name mesa_frames

(F821)


489-489: Undefined name mesa_frames

(F821)


490-490: Undefined name mesa_frames

(F821)


512-512: Undefined name mesa_frames

(F821)


513-513: Undefined name mesa_frames

(F821)


535-535: Undefined name mesa_frames

(F821)


536-536: Undefined name mesa_frames

(F821)


667-667: Undefined name mesa_frames

(F821)


687-687: Undefined name mesa_frames

(F821)


709-709: Undefined name mesa_frames

(F821)


747-747: Undefined name mesa_frames

(F821)


759-759: Undefined name mesa_frames

(F821)


772-772: Undefined name mesa_frames

(F821)


797-797: Undefined name mesa_frames

(F821)


801-801: Undefined name mesa_frames

(F821)

mesa_frames/concrete/mixin.py

54-54: mesa_frames.types_.DataFrame imported but unused

Remove unused import: mesa_frames.types_.DataFrame

(F401)

🪛 GitHub Actions: build
mesa_frames/concrete/agents.py

[error] 51-51: ImportError: cannot import name 'Self' from 'typing'. This causes test collection failures in multiple test modules.

🔇 Additional comments (36)
mesa_frames/__init__.py (2)

44-45: Good addition of postponed evaluation for annotations.

Adding from __future__ import annotations is a good practice for modern Python type checking. This enables postponed evaluation of type annotations, which helps avoid circular import issues and aligns with PEP 563.


48-48: Import reordering looks good.

The reordering of the import for ModelDF from mesa_frames.concrete.model maintains the correctness of imports while potentially resolving dependency ordering requirements.

pyproject.toml (2)

81-81: Good replacement of typeguard with beartype.

Replacing typeguard with beartype for runtime type checking aligns with the PR objectives. Beartype is known for its high-performance runtime type checking with minimal overhead, which should provide better debugging experience as intended.


88-88: Beneficial addition of pre-commit.

Adding pre-commit as a dev dependency is a good practice. It will help enforce consistent code quality and potentially automate type checking before commits, supporting the runtime type checking objectives.

mesa_frames/types_.py (6)

3-13: Good modernization of imports and dependencies.

The additions of from __future__ import annotations and comprehensive imports from standard libraries and beartype improve type checking capabilities. The organization of imports is clean and follows best practices.


15-30: Well-handled optional dependencies.

The implementation of optional imports with fallback dummy classes is an excellent pattern. This ensures that type annotations remain valid even if optional dependencies like pandas or pyarrow aren't installed, improving the package's resilience.


41-55: Enhanced type definitions for Polars types.

The refined type definitions for Polars types are more precise and comprehensive. The union types properly capture the range of acceptable inputs, which will make the runtime type checking more accurate and helpful for developers.


56-70: Comprehensive IntoExpr type definition.

The new IntoExpr type alias is well-defined and covers a wide range of primitive and complex types that can be converted to expressions. This will provide better type checking for functions that accept expression-convertible inputs.


73-84: Improved generic type aliases with special attention to Infinity.

The refined generic type aliases improve type precision. The Infinity type using Annotated with IsEqual[math.inf] is particularly clever - it ensures that only mathematical infinity is accepted where this type is used, preventing potential logical errors.


117-117: Enhanced DiscreteSpaceCapacity type.

Adding np.ndarray to the DiscreteSpaceCapacity union type improves flexibility when working with capacity specifications, allowing direct usage of NumPy arrays without conversion.

.gitignore (1)

155-156: Appropriate additions to .gitignore.

Adding *.code-workspace and llm_rules.md to the ignore list prevents IDE-specific files and presumably LLM-related documentation from being tracked in version control, which is a good practice for keeping the repository clean.

tests/test_agentset.py (1)

5-5: Type checking update to beartype looks good.

The change from typeguard to beartype for runtime type checking aligns with the PR objectives. Beartype is known for being more efficient with lower runtime overhead compared to typeguard.

Also applies to: 11-11

tests/test_grid.py (1)

4-4: Good migration to beartype for runtime type checking.

The decorator change from typeguard.typechecked to beartype.beartype is consistent with the overall migration strategy. This change maintains the runtime type checking functionality while potentially improving performance.

Also applies to: 19-19

tests/test_mixin.py (1)

4-4: Consistent adoption of beartype for TestPolarsMixin.

The change maintains consistency across the codebase by replacing typeguard with beartype. This is aligned with the project's overall direction to use beartype for runtime type validation.

Also applies to: 9-9

tests/test_agents.py (1)

510-510: Improved Series creation approach.

Changing from implicit column creation to using pl.Series() explicitly is a more type-safe approach. This change likely helps beartype better understand and validate the types at runtime.

mesa_frames/concrete/model.py (5)

43-43: Good addition of future imports

Adding from __future__ import annotations enables postponed evaluation of type annotations, allowing for self-referential and forward references in type hints.


48-50: Clean import organization

The imports are well-organized with beartype added for runtime type checking and AgentSetDF moved from conditional to direct import to support runtime type validation.


55-55: Runtime type checking enabled with beartype

Adding the @beartype decorator to the ModelDF class enables runtime type checking, which will help catch type-related errors early.


104-104: Improved return type annotation

Changing the return type from string literal "AgentSetDF" to a direct reference AgentSetDF improves type checking and IDE support.


222-227: Good addition of docstring

Adding a docstring to the space setter method improves code documentation and helps users understand the parameter requirements.

mesa_frames/concrete/agentset.py (2)

60-72: Well-organized imports

The import organization is good with future annotations for postponed evaluation, beartype for runtime type checking, and clear typing imports. The removal of conditional TYPE_CHECKING imports makes types available at runtime.

🧰 Tools
🪛 Ruff (0.8.2)

67-67: typing.Literal imported but unused

Remove unused import: typing.Literal

(F401)


71-71: mesa_frames.concrete.model.ModelDF imported but unused

Remove unused import: mesa_frames.concrete.model.ModelDF

(F401)


72-72: mesa_frames.types_.IdsLike imported but unused

Remove unused import: mesa_frames.types_.IdsLike

(F401)


76-77: Runtime type checking with beartype

The @beartype decorator properly implements runtime type checking while maintaining the class documentation with @copydoc.

mesa_frames/concrete/space.py (5)

45-45: Good import organization

The imports are well organized with direct import of inf from math, beartype for runtime type checking, and the new Infinity type from mesa_frames.types_. This supports both the runtime type checking objective and improved type safety.

Also applies to: 51-51, 55-55


59-59: Runtime type checking enabled with beartype

Adding the @beartype decorator to the GridPolars class enables runtime type checking, which will help catch type-related errors early.


89-89: More flexible parameter type

Changing the capacity parameter type from int to int | None makes the API more flexible, allowing for optional capacity specification.


168-168: Enhanced input type flexibility

Changing the agents parameter type from pl.DataFrame to pl.DataFrame | pl.Series provides more flexibility for function callers.


219-222: Better infinity handling

The changes to return type and implementation improve infinity handling by:

  1. Using a more precise return type with Infinity
  2. Using Python's built-in inf instead of NumPy's np.inf
  3. Casting the sum to an integer for consistent return types when not infinite
mesa_frames/abstract/mixin.py (4)

45-45: Well-organized imports

The import organization is good with future annotations for postponed evaluation, beartype for runtime type checking, and consolidated type imports including the new DataFrameInput type.

Also applies to: 52-62


65-65: Runtime type checking with beartype

Adding the @beartype decorator to both CopyMixin and DataFrameMixin classes provides runtime type checking which helps catch type errors early.

Also applies to: 187-187


272-272: More flexible type annotations

Changing specific collection types to more general Collection types allows for greater flexibility while maintaining type safety.

Also applies to: 478-480


343-343: Enhanced type flexibility

Multiple parameter types have been updated to accept more general collections or broader ranges of input types, providing greater flexibility for function callers.

Also applies to: 467-467, 487-487, 496-497, 505-506

mesa_frames/concrete/agents.py (2)

75-83: Avoid hard dependency on fully-qualified forward reference at runtime

beartype eagerly evaluates annotations; therefore
mesa_frames.concrete.model.ModelDF must resolve at import time, risking
circular-import failures.
After adding the above TYPE_CHECKING guard you can safely shorten the hint:

- def __init__(self, model: mesa_frames.concrete.model.ModelDF) -> None:
+ def __init__(self, model: ModelDF) -> None:

Keep the long path only inside the TYPE_CHECKING block.
[ suggest_essential_refactor ]

🧰 Tools
🪛 Ruff (0.8.2)

75-75: Undefined name mesa_frames

(F821)


403-411: Minor: rename loop variable to prevent shadowing but use it consistently

Great that you renamed maskmask_value to avoid confusion, yet inside
return_dictionary[agentset] = … you’re already using mask_value, so the
comment above (“change inner loop variable”) can be deleted to reduce clutter.
[ suggest_nitpick ]

mesa_frames/concrete/mixin.py (1)

343-355: Good use of overloads to improve _df_index typing

Splitting scalar vs. multi-column overloads clarifies call sites and helps
static analysers. 👍

mesa_frames/abstract/agents.py (2)

829-845: discard() delegates correctly—nice reuse

Delegating to the parent implementation keeps behaviour consistent while
tightening the signature.
Looks good.


1129-1132: Potential NoneType access when space is not set

self.space may be None until the model sets it.
self.space.agents will then raise AttributeError.

Consider:

if self.space is None:
    raise ValueError("Space not initialised for this model.")

before accessing it, or document the pre-condition.

@adamamer20
Copy link
Member Author

@coderabbitai review

Copy link
Contributor

coderabbitai bot commented Apr 29, 2025

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🔭 Outside diff range comments (2)
mesa_frames/abstract/agents.py (2)

43-52: ⚠️ Potential issue

Missing imports – unresolved names break beartype at import-time

@beartype resolves annotations immediately when the decorator is executed.
Because mesa_frames and its sub-modules are never imported in this file, every
fully-qualified annotation such as mesa_frames.concrete.model.ModelDF is
currently a dangling forward reference. This triggers a NameError during
module import and is also the root-cause of the Ruff F821 warnings.

-from abc import abstractmethod
-from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
+from abc import abstractmethod
+from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
 from contextlib import suppress
 from beartype import beartype
 from numpy.random import Generator
 from typing import Self, overload, Any, Literal
+
+# --------------------------------------------------------------------------- #
+# Imports required only for annotations but *must* exist at runtime so that   #
+# `beartype` can resolve them via `typing.get_type_hints`.  Using `# noqa`    #
+# silences the “unused import” warnings without disabling strict type checks. #
+# --------------------------------------------------------------------------- #
+from mesa_frames.concrete.model import ModelDF  # noqa: F401
+from mesa_frames.concrete.agents import AgentSetDF  # noqa: F401
+from mesa_frames.abstract.space import SpaceDF  # noqa: F401

This keeps the readable, fully-qualified annotations while guaranteeing that
names exist in the module namespace for both static analysis and runtime
inspection. (If circular-import concerns arise, move these imports inside a
if TYPE_CHECKING: block and convert annotations back to strings, but that
negates the stated PR goal of avoiding TYPE_CHECKING blocks.)


271-299: ⚠️ Potential issue

select() lost its @abstractmethod decorator – subclasses get a stub that returns ...

The method body is only ..., yet the decorator was removed. As a result:

  • AgentContainer is no longer abstract with respect to select.
  • Any concrete subclass that does not override select will inherit a
    method that returns an Ellipsis, leading to hard-to-trace runtime errors.
-    def select(
+    @abstractmethod
+    def select(

Either re-apply @abstractmethod or provide a concrete, working implementation.

🧹 Nitpick comments (1)
mesa_frames/abstract/agents.py (1)

932-945: Edge-case in remove() – magic string & duplicated traversal

  1. Overloading agents="active" mixes ID semantics with a command keyword and
    can silently clash with a legitimate agent whose unique_id is the literal
    string "active". Consider an explicit boolean flag instead, e.g.
    remove_active: bool = False.

  2. The loop that searches the returned agentsdf for an AgentSetDF of the
    same class is O(n) per call. If agentsdf.agents were keyed by the
    agent-set name (or type) the lookup could be constant-time.

🧰 Tools
🪛 Ruff (0.8.2)

941-941: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between de7da65 and e183449.

⛔ Files ignored due to path filters (1)
  • uv.lock is excluded by !**/*.lock
📒 Files selected for processing (5)
  • .github/workflows/build.yml (1 hunks)
  • .pre-commit-config.yaml (1 hunks)
  • mesa_frames/abstract/agents.py (27 hunks)
  • mesa_frames/concrete/mixin.py (15 hunks)
  • pyproject.toml (2 hunks)
✅ Files skipped from review due to trivial changes (1)
  • .github/workflows/build.yml
🚧 Files skipped from review as they are similar to previous changes (2)
  • pyproject.toml
  • mesa_frames/concrete/mixin.py
🧰 Additional context used
🪛 Ruff (0.8.2)
mesa_frames/abstract/agents.py

71-71: Undefined name mesa_frames

(F821)


80-80: Undefined name mesa_frames

(F821)


81-81: Undefined name mesa_frames

(F821)


107-107: Undefined name mesa_frames

(F821)


108-108: Undefined name mesa_frames

(F821)


134-134: Undefined name mesa_frames

(F821)


139-139: Undefined name mesa_frames

(F821)


176-176: Undefined name mesa_frames

(F821)


187-187: Undefined name mesa_frames

(F821)


250-250: Undefined name mesa_frames

(F821)


251-251: Undefined name mesa_frames

(F821)


399-399: Undefined name mesa_frames

(F821)


400-400: Undefined name mesa_frames

(F821)


489-489: Undefined name mesa_frames

(F821)


490-490: Undefined name mesa_frames

(F821)


512-512: Undefined name mesa_frames

(F821)


513-513: Undefined name mesa_frames

(F821)


535-535: Undefined name mesa_frames

(F821)


536-536: Undefined name mesa_frames

(F821)


667-667: Undefined name mesa_frames

(F821)


687-687: Undefined name mesa_frames

(F821)


709-709: Undefined name mesa_frames

(F821)


747-747: Undefined name mesa_frames

(F821)


759-759: Undefined name mesa_frames

(F821)


772-772: Undefined name mesa_frames

(F821)


797-797: Undefined name mesa_frames

(F821)


801-801: Undefined name mesa_frames

(F821)

🔇 Additional comments (1)
.pre-commit-config.yaml (1)

20-20: Ensure pyupgrade supports the --py311-plus flag
The --py311-plus argument is only available in recent pyupgrade releases. Please verify that the pinned version (v3.19.1) indeed supports this flag to prevent pre-commit failures.

Copy link

codecov bot commented Apr 29, 2025

Codecov Report

Attention: Patch coverage is 94.21488% with 7 lines in your changes missing coverage. Please review.

Project coverage is 91.98%. Comparing base (8c8fb4b) to head (b880ecd).
Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
mesa_frames/types_.py 82.14% 5 Missing ⚠️
mesa_frames/abstract/agents.py 92.59% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #143      +/-   ##
==========================================
+ Coverage   91.78%   91.98%   +0.19%     
==========================================
  Files          11       11              
  Lines        1680     1721      +41     
==========================================
+ Hits         1542     1583      +41     
  Misses        138      138              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

adamamer20 and others added 2 commits April 29, 2025 13:42
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
@adamamer20 adamamer20 requested review from EwoutH and Ben-geo April 29, 2025 11:48
@Ben-geo
Copy link
Collaborator

Ben-geo commented May 2, 2025

Would it make sense to structure this so that runtime type checking is enabled only in development or test environments? For example, we could use a wrapper like:

def implement_beartype(func):
    if os.getenv("ENV_TYPE") in ["dev", "test"]:
        return beartype(func)
    return func

This way we could annotate functions with @implement_beartype, and control activation via environment variables. It would also allow more advanced users to opt-in to runtime type checking in production if they really want to

i guess this will stop the over head from affecting everyone

@adamamer20
Copy link
Member Author

Yes @Ben-geo, I think the best thing would be to gate all runtime type-checking behind a single env var: MESA_FRAMES_RUNTIME_TYPECHECKING. In our top-level __init__.py we’d do:

import os
from beartype.claw import beartype_this_package

if os.getenv("MESA_FRAMES_RUNTIME_TYPECHECKING", "").lower() in ("1", "true", "yes"):
    beartype_this_package()
  1. Tests (CI): set MESA_FRAMES_RUNTIME_TYPECHECKING=true so every import is wrapped and any signature mismatch fails the build.

    export MESA_FRAMES_RUNTIME_TYPECHECKING=true
    pytest
  2. Local dev: opt in by exporting the same var when you want real-time feedback.

    export MESA_FRAMES_RUNTIME_TYPECHECKING=1
    python -c "import mesa_frames; print('Beartype active!')"

    Or programmatically:

    import os
    os.environ["MESA_FRAMES_RUNTIME_TYPECHECKING"] = "true"
    import mesa_frames  # now beartype checks will be enabled
  3. Production: leave the var unset (or run under -O) and pay zero cost—no wrappers get installed.

This keeps our codebase clean, gives us type-safety in CI/dev, and preserves peak performance everywhere else.

@Ben-geo
Copy link
Collaborator

Ben-geo commented May 3, 2025

Awesome @adamamer20 I think this is the best solution.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
feature New functionality added to the project.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Run-time type checking
2 participants