From e8784db67604343f5377f396547196c862d71b14 Mon Sep 17 00:00:00 2001
From: ahalev <avishaihalev@gmail.com>
Date: Thu, 30 Mar 2023 17:49:30 -0700
Subject: [PATCH 1/2] add dry_run docstring

---
 src/pymgrid/utils/dry_run.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/src/pymgrid/utils/dry_run.py b/src/pymgrid/utils/dry_run.py
index c8f27c96..b6577be7 100644
--- a/src/pymgrid/utils/dry_run.py
+++ b/src/pymgrid/utils/dry_run.py
@@ -8,6 +8,31 @@
 
 @contextmanager
 def dry_run(pymgrid_object: Union['pymgrid.Microgrid', 'pymgrid.modules.base.BaseMicrogridModule']):
+    """
+    A context manager to test modifications of a pymgrid object without modifying said object.
+
+    Parameters
+    ----------
+    pymgrid_object : :class:`pymgrid.Microgrid` or :class:`pymgrid.modules.base.BaseMicrogridModule`
+        Microgrid or module to try run
+
+    Returns
+    -------
+    test_object : copy of object to test usage of, without modifying original object.
+
+    Examples
+    --------
+    >>> from pymgrid.modules import BatteryModule
+    >>> module = BatteryModule(0, 100, 50, 50, 0.9, init_soc=0.5)
+    >>> with dry_run(module) as test_module:
+    >>>     test_module.step(test_module.max_act, normalized=False)
+    >>>     print(f'Current step: {test_module.current_step}; current charge: {test_module.current_charge}')
+    Current step: 1; current charge: 0.0
+    >>> print(f'Current step: {module.current_step}; current charge: {module.current_charge}')
+    Current step: 0; current charge: 50.0
+
+    """
+
     serialized = yaml.safe_dump(pymgrid_object)
 
     try:

From 16fdf58c9db35c27c7deaf44801117f945a0dc3a Mon Sep 17 00:00:00 2001
From: ahalev <avishaihalev@gmail.com>
Date: Thu, 30 Mar 2023 17:49:59 -0700
Subject: [PATCH 2/2] remove type hints

---
 src/pymgrid/utils/dry_run.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/pymgrid/utils/dry_run.py b/src/pymgrid/utils/dry_run.py
index b6577be7..872e43bf 100644
--- a/src/pymgrid/utils/dry_run.py
+++ b/src/pymgrid/utils/dry_run.py
@@ -1,13 +1,10 @@
 import yaml
 
 from contextlib import contextmanager
-from typing import Union
-
-import pymgrid
 
 
 @contextmanager
-def dry_run(pymgrid_object: Union['pymgrid.Microgrid', 'pymgrid.modules.base.BaseMicrogridModule']):
+def dry_run(pymgrid_object):
     """
     A context manager to test modifications of a pymgrid object without modifying said object.