Add and make use of small python library for yarns

Due to an unwillingness to add another IMPLEMENTS that copypasta'd
the same "MATCH_n = os.environ['MATCH_n']", add this small library
in an attempt to reduce the amount of repitition.

Change-Id: I64dc67ad7bd4c0d7572906168c72b0628a7574db

2 files changed, 81 insertions, 58 deletions

diff --git a/yarns.webapp/900-implementations.yarn b/yarns.webapp/900-implementations.yarn
index 0fec328..6e0b380 100644
--- a/yarns.webapp/900-implementations.yarn
+++ b/yarns.webapp/900-implementations.yarn
@@ -85,15 +85,12 @@ Add a `lorries` section to a `lorry-controller.conf`. This hardcodes
 most of the configuration.
 
     IMPLEMENTS GIVEN (\S+) in (\S+) adds lorries (\S+) using prefix (\S+)
+    cd "$SRCDIR"/yarns.webapp
     python -c '
     import os
-    import json
+    import yarnlib
     
-    DATADIR = os.environ["DATADIR"]
-    MATCH_1 = os.environ["MATCH_1"]
-    MATCH_2 = os.environ["MATCH_2"]
-    MATCH_3 = os.environ["MATCH_3"]
-    MATCH_4 = os.environ["MATCH_4"]
+    MATCH_1, MATCH_2, MATCH_3, MATCH_4 = yarnlib.matches()
     
     new = {
         "type": "lorries",
@@ -104,26 +101,22 @@ most of the configuration.
         ],
     }
 
-    filename = os.path.join(DATADIR, MATCH_2, MATCH_1)
-    with open(filename, "r") as f:
-        obj = json.load(f)
+    filename = os.path.join(yarnlib.DATADIR, MATCH_2, MATCH_1)
+    obj = yarnlib.load_json_from_file(filename)
     obj.append(new)
-    with open(filename, "w") as f:
-        json.dump(obj, f)
+    yarnlib.dump_json_to_file(obj, filename)
     '
 
 Add a `troves` section to `lorry-controller.conf`. Again, we hardcode
 most of the configuration.
 
     IMPLEMENTS GIVEN (\S+) in (\S+) adds trove (\S+)
+    cd "$SRCDIR"/yarns.webapp
     python -c '
     import os
-    import json
-    
-    DATADIR = os.environ["DATADIR"]
-    MATCH_1 = os.environ["MATCH_1"]
-    MATCH_2 = os.environ["MATCH_2"]
-    MATCH_3 = os.environ["MATCH_3"]
+    import yarnlib
+
+    MATCH_1, MATCH_2, MATCH_3 = yarnlib.matches()
     
     new = {
         "type": "troves",
@@ -134,93 +127,76 @@ most of the configuration.
         "prefixmap": {},
     }
 
-    filename = os.path.join(DATADIR, MATCH_2, MATCH_1)
-    with open(filename, "r") as f:
-        obj = json.load(f)
+    filename = os.path.join(yarnlib.DATADIR, MATCH_2, MATCH_1)
+    obj = yarnlib.load_json_from_file(filename)
     obj.append(new)
-    with open(filename, "w") as f:
-        json.dump(obj, f, indent=4)
+    yarnlib.dump_json_to_file(obj, filename)
     '
 
 Set the a specific field for all sections in a `lorry-controller.conf`
 file.
 
     IMPLEMENTS GIVEN (\S+) in (\S+) has (\S+) set to (.+) for everything
+    cd "$SRCDIR"/yarns.webapp
     python -c '
     import os
     import json
+    import yarnlib
     
-    DATADIR = os.environ["DATADIR"]
-    MATCH_1 = os.environ["MATCH_1"]
-    MATCH_2 = os.environ["MATCH_2"]
-    MATCH_3 = os.environ["MATCH_3"]
-    MATCH_4 = os.environ["MATCH_4"]
-    
-    filename = os.path.join(DATADIR, MATCH_2, MATCH_1)
+    MATCH_1, MATCH_2, MATCH_3, MATCH_4 = yarnlib.matches()
+
+    filename = os.path.join(yarnlib.DATADIR, MATCH_2, MATCH_1)
 
-    with open(filename, "r") as f:
-        obj = json.load(f)
+    obj = yarnlib.load_json_from_file(filename)
 
     for section in obj:
         section[MATCH_3] = json.loads(MATCH_4)
 
-    with open(filename, "w") as f:
-        json.dump(obj, f, indent=4)
+    yarnlib.dump_json_to_file(obj, filename)
     '
 
 Set a specific field for a `troves` section.
 
     IMPLEMENTS GIVEN (\S+) in (\S+) sets (\S+) to (.+) for trove (\S+)
+    cd "$SRCDIR"/yarns.webapp
     python -c '
     import os
     import json
-    
-    DATADIR = os.environ["DATADIR"]
-    MATCH_1 = os.environ["MATCH_1"]
-    MATCH_2 = os.environ["MATCH_2"]
-    MATCH_3 = os.environ["MATCH_3"]
-    MATCH_4 = os.environ["MATCH_4"]
-    MATCH_5 = os.environ["MATCH_5"]
-    
-    filename = os.path.join(DATADIR, MATCH_2, MATCH_1)
+    import yarnlib
+
+    MATCH_1, MATCH_2, MATCH_3, MATCH_4, MATCH_5 = yarnlib.matches()
 
-    with open(filename, "r") as f:
-        obj = json.load(f)
+    filename = os.path.join(yarnlib.DATADIR, MATCH_2, MATCH_1)
+
+    obj = yarnlib.load_json_from_file(filename)
 
     for section in obj:
         if section["type"] in ["trove", "troves"]:
             if section["trovehost"] == MATCH_5:
                 section[MATCH_3] = json.loads(MATCH_4)
 
-    with open(filename, "w") as f:
-        json.dump(obj, f, indent=4)
+    yarnlib.dump_json_to_file(obj, filename)
     '
 
 Set the prefixmap for a Trove in a Lorry Controller configuration
 file. Note that the Trove must already be in the configuration file.
 
     IMPLEMENTS GIVEN (\S+) in (\S+) has prefixmap (\S+):(\S+) for (\S+)
+    cd "$SRCDIR"/yarns.webapp
     python -c '
     import os
-    import json
+    import yarnlib
     
-    DATADIR = os.environ["DATADIR"]
-    MATCH_1 = os.environ["MATCH_1"]
-    MATCH_2 = os.environ["MATCH_2"]
-    MATCH_3 = os.environ["MATCH_3"]
-    MATCH_4 = os.environ["MATCH_4"]
-    MATCH_5 = os.environ["MATCH_5"]
+    MATCH_1, MATCH_2, MATCH_3, MATCH_4, MATCH_5 = yarnlib.matches()
     
-    filename = os.path.join(DATADIR, MATCH_2, MATCH_1)
-    with open(filename, "r") as f:
-        objs = json.load(f)
+    filename = os.path.join(yarnlib.DATADIR, MATCH_2, MATCH_1)
+    objs = yarnlib.load_json_from_file(filename)
 
     for obj in objs:
         if obj["type"] == "troves" and obj["trovehost"] == MATCH_5:
             obj["prefixmap"][MATCH_3] = MATCH_4
 
-    with open(filename, "w") as f:
-        json.dump(objs, f, indent=4)
+    yarnlib.dump_json_to_file(objs, filename)
     '
 
 We need to be able to tell WEBAPP, when it runs, where the
diff --git a/yarns.webapp/yarnlib.py b/yarns.webapp/yarnlib.py
new file mode 100644
index 0000000..89f87d9
--- /dev/null
+++ b/yarns.webapp/yarnlib.py
@@ -0,0 +1,47 @@
+# Copyright (C) 2015 Codethink Limited
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+# Python library for lorry-controller yarns.
+
+import os
+import json
+
+DATADIR = os.environ['DATADIR']
+
+def matches():
+    '''Returns a generator of MATCH_n environment variables.'''
+    matches = {}
+    prefix = "MATCH_"
+
+    for key, value in os.environ.iteritems():
+        if key.startswith(prefix):
+            stripped_key = key[len(prefix):]
+            if stripped_key.isdigit() and stripped_key != "0":
+                matches[int(stripped_key)] = value
+
+    for key in sorted(matches):
+        if key != 1 and key-1 not in matches:
+            break
+        yield matches[key]
+
+def load_json_from_file(filename):
+    '''Deserialise json file.'''
+    with open(filename, "r") as f:
+        return json.load(f)
+
+def dump_json_to_file(j, filename):
+    '''Write serialised object to file.'''
+    with open(filename, "w") as f:
+        return json.dump(j, f, indent=4)