PolicyEngine / policyengine-core

Core microsimulation engine for PolicyEngine models. Forked from OpenFisca-Core.
https://policyengine.github.io/policyengine-core
GNU Affero General Public License v3.0
14 stars 19 forks source link

Refactor and optimize Enum class #188

Closed MaxGhenis closed 4 months ago

MaxGhenis commented 4 months ago

Fixes #100

This pull request focuses on refactoring and optimizing the Enum class in the enum.py module. The main goals of this PR are to improve code clarity, performance, and maintainability.

Key changes:

  1. Added docstrings:

    • Added a docstring to the __init__ method to explain its purpose and how the index is assigned to each Enum item.
    • Improved the docstring of the encode method to provide a clear explanation of its functionality, arguments, return value, and examples of usage.
  2. Simplified method assignments:

    • Simplified the __eq__ and __hash__ method assignments by directly assigning them to the corresponding methods of object.
  3. Optimized encode method:

    • Consolidated the .astype(ENUM_ARRAY_DTYPE) call at the end of the method, applying it to the indices array before creating the EnumArray.
    • Simplified the code for handling enum items arrays by using np.select() instead of separate branches for bytes and non-bytes cases.
    • Raised a ValueError with a clear error message for unsupported array data types.
  4. Updated type hints:

    • Updated the type hint for the array parameter in the encode method to use np.ndarray instead of specific data types like np.int_, np.float_, and np.object_.
  5. Improved comments:

    • Updated the comments to provide more context and explain the purpose of each code block.
    • Removed unnecessary comments and simplified existing ones to enhance code readability.
  6. Removed unnecessary code:

    • Removed the warnings import and the warnings.simplefilter() call since they are not directly related to the functionality of the Enum class, and other changes should remove the warnings.

These improvements aim to make the Enum class more readable, efficient, and maintainable. The added docstrings and comments provide better documentation and understanding of the code, while the optimized operations and simplified code structure enhance its performance and clarity.

The consolidation of the .astype(ENUM_ARRAY_DTYPE) call at the end of the encode method reduces code duplication and improves readability. The simplified code for handling enum items arrays using np.select() makes the code more concise and easier to understand.

Please review the changes and provide any feedback or suggestions for further improvements.