From b655f2d17824f2fdb794e49f1023e7720fa06a0c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20Kvalsvik?= <jokva@statoil.com>
Date: Fri, 16 Dec 2016 16:38:24 +0100
Subject: [PATCH] Get groups

---
 python/sunbeam/sunbeam.cpp | 14 ++++++++++++--
 python/sunbeam/sunbeam.py  | 14 ++++++++++++++
 tests/schedule.py          |  7 +++++++
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/python/sunbeam/sunbeam.cpp b/python/sunbeam/sunbeam.cpp
index 058dc0e4c..1bb96b7cd 100644
--- a/python/sunbeam/sunbeam.cpp
+++ b/python/sunbeam/sunbeam.cpp
@@ -25,7 +25,7 @@ struct ptime_to_python_datetime {
 };
 
 template< typename T >
-py::list vector_to_pylist( const std::vector< T >& v ) {
+py::list iterable_to_pylist( const T& v ) {
     py::list l;
     for( const auto& x : v ) l.append( x );
     return l;
@@ -39,6 +39,10 @@ std::vector< Well > get_wells( const Schedule& sch ) {
     return wells;
 }
 
+py::list group_wellnames( const Group& g, size_t timestep ) {
+    return iterable_to_pylist( g.getWells( timestep )  );
+}
+
 /* alias some of boost's long names and operations */
 using ref = py::return_internal_reference<>;
 using copy = py::return_value_policy< py::copy_const_reference >;
@@ -81,7 +85,7 @@ py::list get_timesteps( const Schedule* s ) {
 
     for( size_t i = 0; i < tm.size(); ++i ) v.push_back( tm[ i ] );
 
-    return vector_to_pylist( v );
+    return iterable_to_pylist( v );
 }
 
 }
@@ -131,12 +135,18 @@ py::class_< std::vector< Well > >( "WellList", py::no_init )
     .def( py::vector_indexing_suite< std::vector< Well > >() )
     ;
 
+py::class_< Group >( "Group", py::no_init )
+    .def( "name",       mkcopy( &Group::name ) )
+    .def( "_wellnames", group_wellnames )
+    ;
+
 py::class_< Schedule >( "Schedule", py::no_init )
     .add_property( "_wells", get_wells )
     .add_property( "start",  get_start_time )
     .add_property( "end",    get_end_time )
     .add_property( "timesteps", get_timesteps )
     .def( "__contains__", &Schedule::hasWell )
+    .def( "_group", &Schedule::getGroup, ref() )
     ;
 
 void (ParseContext::*ctx_update)(const std::string&, InputError::Action) = &ParseContext::update;
diff --git a/python/sunbeam/sunbeam.py b/python/sunbeam/sunbeam.py
index 5dd22a7f4..0a4667d84 100644
--- a/python/sunbeam/sunbeam.py
+++ b/python/sunbeam/sunbeam.py
@@ -46,6 +46,9 @@ class Schedule(object):
     def wells(self):
         return map(Well, self._wells)
 
+    def group(self, name):
+        return Group(self._group(name), self)
+
 @delegate(lib.EclipseState)
 class EclipseState(object):
     @property
@@ -75,6 +78,17 @@ class Well(object):
         def fn(well): return well.isproducer(timestep)
         return fn
 
+@delegate(lib.Group)
+class Group(object):
+    def __init__(self, _, schedule):
+        self._schedule = schedule
+
+    def wells(self, timestep):
+        names = self._wellnames(timestep)
+        wells = { well.name: well for well in self._schedule.wells }
+        return map(wells.__getitem__, filter(wells.__contains__, names))
+
+
 def _parse_context(actions):
     ctx = lib.ParseContext()
 
diff --git a/tests/schedule.py b/tests/schedule.py
index 575504991..3d14458ac 100644
--- a/tests/schedule.py
+++ b/tests/schedule.py
@@ -25,3 +25,10 @@ class TestSchedule(unittest.TestCase):
         timesteps = self.sch.timesteps
         self.assertEqual(176, len(timesteps))
         self.assertEqual(dt.date(2015, 12, 31), timesteps[7])
+
+    def testGroups(self):
+        g1 = self.sch.group('G1').wells(0)
+        self.assertEqual(2, len(g1))
+
+        with self.assertRaises(ValueError):
+            self.sch.group('foo')