Well status functions

python/sunbeam/sunbeam.py

@@ -84,6 +84,25 @@ class Well(object):
         def fn(well): return well.isproducer(timestep)
         return fn
 
+    # using the names flowing and closed for functions that test if a well is
+    # opened or closed at some point, because we might want to use the more
+    # imperative words 'open' and 'close' (or 'shut') for *changing* the status
+    # later
+    @staticmethod
+    def flowing(timestep):
+        def fn(well): return well.status(timestep) == 'OPEN'
+        return fn
+
+    @staticmethod
+    def closed(timestep):
+        def fn(well): return well.status(timestep) == 'SHUT'
+        return fn
+
+    @staticmethod
+    def auto(timestep):
+        def fn(well): return well.status(timestep) == 'AUTO'
+        return fn
+
 @delegate(lib.Group)
 class Group(object):
     def __init__(self, _, schedule):

tests/wells.py

@@ -73,3 +73,15 @@ class TestWells(unittest.TestCase):
 
         self.assertEqual(inje[0].name, "INJ")
         self.assertEqual(prod[0].name, "PROD")
+
+    def testOpenFilter(self):
+        open_at_1 = sunbeam.Well.flowing(1)
+        flowing = list(filter(open_at_1, self.wells))
+        closed  = list(filter(lambda well: not open_at_1(well), self.wells))
+
+        self.assertEqual(2, len(flowing))
+        self.assertEqual(0, len(closed))
+
+        flowing1 = filter(lambda x: not sunbeam.Well.closed(1)(x), self.wells)
+        closed1  = filter(sunbeam.Well.closed(1), self.wells)
+        self.assertListEqual(list(closed), list(closed1))