__radd__ in Borehole

Author: MassimoCimmino

pygfunction/boreholes.py

@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+from typing import Union
 import warnings
 
 import numpy as np
@@ -55,11 +56,11 @@ def __repr__(self):
              f' orientation={self.orientation})')
         return s
 
-    def __add__(self, other: Self):
+    def __add__(self, other: Union[Self, list]):
         """
         Adds two boreholes together to form a borefield
         """
-        if not isinstance(other, self.__class__):
+        if not isinstance(other, (self.__class__, list)):
             # Check if other is a borefield and try the operation using
             # other.__radd__
             try:
@@ -70,12 +71,41 @@ def __add__(self, other: Self):
                     f'Expected Borefield, list or Borehole input;'
                     f' got {other}'
                     )
+        elif isinstance(other, list):
+            # Create a borefield from the borehole and a list
+            from .borefield import Borefield
+            field = Borefield.from_boreholes([self] + other)
         else:
             # Create a borefield from the two boreholes
             from .borefield import Borefield
             field = Borefield.from_boreholes([self, other])
         return field
 
+    def __radd__(self, other: Union[Self, list]):
+        """
+        Adds two boreholes together to form a borefield
+        """
+        if not isinstance(other, (self.__class__, list)):
+            # Check if other is a borefield and try the operation using
+            # other.__radd__
+            try:
+                field = other.__add__(self)
+            except:
+                # Invalid input
+                raise TypeError(
+                    f'Expected Borefield, list or Borehole input;'
+                    f' got {other}'
+                    )
+        elif isinstance(other, list):
+            # Create a borefield from the borehole and a list
+            from .borefield import Borefield
+            field = Borefield.from_boreholes(other + [self])
+        else:
+            # Create a borefield from the two boreholes
+            from .borefield import Borefield
+            field = Borefield.from_boreholes([other, self])
+        return field
+
     def distance(self, target):
         """
         Evaluate the distance between the current borehole and a target

tests/borefield_test.py

@@ -32,6 +32,30 @@ def test_borefield_init(field, request):
         H, D, r_b, x, y, tilt=tilt, orientation=orientation)
     assert borefield == borefield_from_boreholes
 
+
+# Test Borefield.__add__ and Borefield.__radd__
+@pytest.mark.parametrize("field, other_field, field_list, other_field_list", [
+        # Using Borefield and Borehole objects
+        ('single_borehole', 'two_boreholes_inclined', False, False),
+        ('ten_boreholes_rectangular', 'single_borehole_short', False, False),
+        ('ten_boreholes_rectangular', 'two_boreholes_inclined', False, False),
+        # Using Borefield as lists
+        ('ten_boreholes_rectangular', 'two_boreholes_inclined', False, True),
+        ('ten_boreholes_rectangular', 'two_boreholes_inclined', True, False),
+    ])
+def test_borefield_add(field, other_field, field_list, other_field_list, request):
+    field = request.field
+    other_field = request.other_field
+    reference_field = gt.borefield.Borefield.from_boreholes(
+        field.to_boreholes() + other_field.to_boreholes()
+        )
+    if field_list:
+        field = field.to_boreholes()
+    if other_field_list:
+        other_field = other_field.to_boreholes()
+    assert field + other_field_list == reference_field
+
+
 # Test borefield comparison using __eq__
 @pytest.mark.parametrize("field, other_field, expected", [
         # Fields that are equal
@@ -53,6 +77,7 @@ def test_borefield_eq(field, other_field, expected, request):
     other_field = request.getfixturevalue(other_field)
     assert (borefield == other_field) == expected
 
+
 # Test borefield comparison using __ne__
 @pytest.mark.parametrize("field, other_field, expected", [
         # Fields that are equal
@@ -75,16 +100,6 @@ def test_borefield_ne(field, other_field, expected, request):
     assert (borefield != other_field) == expected
 
 
-def test_borefield_add():
-    borehole = gt.boreholes.Borehole(100, 1, 0.075, 15, 10)
-    borefield = gt.borefield.Borefield.rectangle_field(2, 1, 6, 6, 100, 1, 0.075)
-    borefield_2 = gt.borefield.Borefield.from_boreholes([borehole, gt.boreholes.Borehole(110, 1, 0.075, 20, 15)])
-    assert borefield + borehole == gt.borefield.Borefield.from_boreholes(borefield.to_boreholes() + [borehole])
-    assert borehole + borefield == gt.borefield.Borefield.from_boreholes(borefield.to_boreholes() + [borehole])
-    assert borefield + [borehole] == gt.borefield.Borefield.from_boreholes(borefield.to_boreholes() + [borehole])
-    assert borefield + borefield_2 == gt.borefield.Borefield.from_boreholes(borefield.to_boreholes()+borefield_2.to_boreholes())
-
-
 # =============================================================================
 # Test evaluate_g_function (vertical boreholes)
 # =============================================================================

tests/boreholes_test.py

@@ -33,6 +33,24 @@ def test_borehole_init():
          ])
 
 
+# Test Borehole.__add__
+@pytest.mark.parametrize("borehole, other_borehole, borehole_list, other_borehole_list", [
+        ('single_borehole', 'single_borehole_short', False, False),
+        ('single_borehole', 'single_borehole_short', True, False),
+        ('single_borehole', 'single_borehole_short', False, True),
+    ])
+def test_borehole_add(borehole, other_borehole, borehole_list, other_borehole_list, request):
+    borehole = request.borehole
+    other_borehole = request.other_borehole
+    field = gt.borefield.Borefield.from_boreholes(
+        [borehole, other_borehole])
+    if borehole_list:
+        borehole = [borehole]
+    if other_borehole_list:
+        other_borehole = [other_borehole]
+    assert field == borehole + other_borehole
+
+
 # Test Borehole.distance
 @pytest.mark.parametrize("borehole1, borehole2", [
         # Same borehole
@@ -307,9 +325,3 @@ def test_circle_field(N, R):
          len(field) == 1 or np.isclose(np.min(dis), B_min),
          len(field) == 1 or np.max(dis) <= (2 + 1e-6) * R,
          ])
-
-
-def test_add_boreholes():
-    borehole1 = gt.boreholes.Borehole(100, 1, 0.075, 0, 0)
-    borehole2 = gt.boreholes.Borehole(110, 1, 0.075, 0, 0)
-    assert gt.borefield.Borefield.from_boreholes([borehole1, borehole2]) == borehole1 + borehole2