Free functions on Well for filtering purposes

python/sunbeam/__init__.py

@@ -1,4 +1,5 @@
 from .sunbeam import parse
+from .sunbeam import Well
 from .libsunbeam import action
 
 __version__     = '0.0.1'

python/sunbeam/sunbeam.cpp

@@ -54,6 +54,9 @@ py::class_< Well >( "Well", py::no_init )
     .def( "J",   headJts )
     .def( "ref", refD )
     .def( "ref", refDts )
+    .def( "isdefined",  &Well::hasBeenDefined )
+    .def( "isinjector", &Well::isInjector )
+    .def( "isproducer", &Well::isProducer )
     ;
 
 py::class_< std::vector< Well > >( "WellList", py::no_init )

python/sunbeam/sunbeam.py

@@ -46,6 +46,24 @@ class EclipseState(object):
     def schedule(self):
         return self._schedule()
 
+@delegate(lib.Well)
+class Well(object):
+
+    @staticmethod
+    def defined(timestep):
+        def fn(well): return well.isdefined(timestep)
+        return fn
+
+    @staticmethod
+    def injector(timestep):
+        def fn(well): return well.isinjector(timestep)
+        return fn
+
+    @staticmethod
+    def producer(timestep):
+        def fn(well): return well.isproducer(timestep)
+        return fn
+
 def _parse_context(actions):
     ctx = lib.ParseContext()
 

tests/wells.py

@@ -22,3 +22,16 @@ class TestWells(unittest.TestCase):
         self.assertEqual(6, i)
         self.assertEqual(6, j)
         self.assertEqual(2247.9, refdepth)
+
+    def testWellDefinedFilter(self):
+        defined0 = filter(sunbeam.Well.defined(0), self.wells)
+        defined1 = filter(sunbeam.Well.defined(1), self.wells)
+        self.assertEqual(len(list(defined0)), 2)
+        self.assertEqual(len(list(defined1)), 2)
+
+    def testWellProdInjeFilter(self):
+        inje = filter(sunbeam.Well.injector(0), self.wells)
+        prod = filter(sunbeam.Well.producer(0), self.wells)
+
+        self.assertEqual(len(list(inje)), 1)
+        self.assertEqual(len(list(prod)), 1)