Skip to content

Commit

Permalink
pre-commit reformatted
Browse files Browse the repository at this point in the history
  • Loading branch information
liuly12 committed Jul 26, 2024
1 parent 1054f65 commit 55a7d05
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 11 deletions.
3 changes: 2 additions & 1 deletion tests/test_distribution.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,8 +69,9 @@ def test_leakage(self):

def test_distribution_overrides(self):
distribution = Distribution(name="", leakage=0.2)
distribution.apply_overrides({'leakage': 0})
distribution.apply_overrides({"leakage": 0})
self.assertEqual(distribution.leakage, 0)


if __name__ == "__main__":
unittest.main()
21 changes: 11 additions & 10 deletions wsimod/nodes/distribution.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,11 @@
@author: bdobson
"""

from typing import Any, Dict

from wsimod.core import constants
from wsimod.nodes.nodes import Node
from typing import Any, Dict


def decorate_leakage_set(self, f):
"""Decorator to extend the functionality of `f` by introducing leakage. This is
Expand Down Expand Up @@ -126,32 +128,31 @@ def __init__(self, leakage=0, **kwargs):
self.push_set_handler["default"] = self.push_set_deny
self.push_check_handler["default"] = self.push_check_deny
self.decorate_pull_handlers()

def decorate_pull_handlers(self):
"""Decorate handlers if there is leakage ratio.
"""
"""Decorate handlers if there is leakage ratio."""
if self.leakage > 0:
self.pull_set_handler["default"] = decorate_leakage_set(
self, self.pull_set_handler["default"]
)
self.pull_check_handler["default"] = decorate_leakage_check(
self, self.pull_check_handler["default"]
)

def apply_overrides(self, overrides: Dict[str, Any] = {}):
"""Apply overrides to the sewer.
Enables a user to override any of the following parameters:
Enables a user to override any of the following parameters:
leakage.
Args:
overrides (dict, optional): Dictionary of overrides. Defaults to {}.
"""
self.leakage = overrides.pop("leakage",
self.leakage)
self.leakage = overrides.pop("leakage", self.leakage)
self.decorate_pull_handlers()
super().apply_overrides(overrides)


class UnlimitedDistribution(Distribution):
""""""

Expand Down

0 comments on commit 55a7d05

Please sign in to comment.