From b649e363001e5089f9483273d1aece82332bff65 Mon Sep 17 00:00:00 2001
From: liuly12 <leyang.liu@hotmail.com>
Date: Sat, 16 Mar 2024 14:22:44 +0000
Subject: [PATCH 1/2] Apply overrides in Distribution

---
 tests/test_distribution.py   |  4 ++++
 wsimod/nodes/distribution.py | 25 +++++++++++++++++++++----
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/tests/test_distribution.py b/tests/test_distribution.py
index aeaabc7..239400e 100644
--- a/tests/test_distribution.py
+++ b/tests/test_distribution.py
@@ -67,6 +67,10 @@ def test_leakage(self):
         self.assertEqual(v2, arc1.vqip_in["volume"])
         self.assertEqual(v2 * 0.2, arc2.vqip_in["volume"])
 
+    def test_distribution_overrides(self):
+        distribution = Distribution(name="", leakage=0.2)
+        distribution.apply_overrides({'leakage': 0})
+        self.assertEqual(distribution.leakage, 0)
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/wsimod/nodes/distribution.py b/wsimod/nodes/distribution.py
index 63a7cc0..fece94d 100644
--- a/wsimod/nodes/distribution.py
+++ b/wsimod/nodes/distribution.py
@@ -6,7 +6,7 @@
 
 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
@@ -125,15 +125,32 @@ def __init__(self, leakage=0, **kwargs):
         # Update handlers
         self.push_set_handler["default"] = self.push_set_deny
         self.push_check_handler["default"] = self.push_check_deny
-
-        if leakage > 0:
+        self.decorate_pull_handlers()
+    
+    def decorate_pull_handlers(self):
+        """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: 
+        leakage.
+        
+        Args:
+            overrides (dict, optional): Dictionary of overrides. Defaults to {}.
+        """
+        self.leakage = overrides.pop("leakage", 
+                                      self.leakage)
+        self.decorate_pull_handlers()
+        super().apply_overrides(overrides)
 
 class UnlimitedDistribution(Distribution):
     """"""

From 55a7d05a63a6ce7b00bfd2e807ded1fea36d9539 Mon Sep 17 00:00:00 2001
From: liuly12 <leyang.liu@hotmail.com>
Date: Fri, 26 Jul 2024 13:24:12 +0100
Subject: [PATCH 2/2] pre-commit reformatted

---
 tests/test_distribution.py   |  3 ++-
 wsimod/nodes/distribution.py | 21 +++++++++++----------
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/tests/test_distribution.py b/tests/test_distribution.py
index 239400e..fca95a2 100644
--- a/tests/test_distribution.py
+++ b/tests/test_distribution.py
@@ -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()
diff --git a/wsimod/nodes/distribution.py b/wsimod/nodes/distribution.py
index fece94d..66baff4 100644
--- a/wsimod/nodes/distribution.py
+++ b/wsimod/nodes/distribution.py
@@ -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
@@ -126,10 +128,9 @@ 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"]
@@ -137,21 +138,21 @@ def decorate_pull_handlers(self):
             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):
     """"""