Added 'block info' to retrieve and export portions of Curve.data[*].info.

[hooke.git] / hooke / plugin / curve.py

diff --git a/hooke/plugin/curve.py b/hooke/plugin/curve.py
index b287ee4fa11bd5ad96f10de2269c6523ebd81bad..92e5bd1bb8a1aeba4b5ca90597c6f685d154fe0d 100644 (file)
--- a/hooke/plugin/curve.py
+++ b/hooke/plugin/curve.py
@@ -25,8 +25,10 @@ associated :class:`hooke.command.Command`\s for handling
 """
 
 import copy
+import re
 
 import numpy
+import yaml
 
 from ..command import Command, Argument, Failure
 from ..command_stack import CommandStack
@@ -263,8 +265,8 @@ class CurvePlugin (Builtin):
     def __init__(self):
         super(CurvePlugin, self).__init__(name='curve')
         self._commands = [
-            GetCommand(self), InfoCommand(self), DeltaCommand(self),
-            ExportCommand(self), DifferenceCommand(self),
+            GetCommand(self), InfoCommand(self), BlockInfoCommand(self),
+            DeltaCommand(self), ExportCommand(self), DifferenceCommand(self),
             DerivativeCommand(self), PowerSpectrumCommand(self),
             ClearStackCommand(self)]
 
@@ -345,6 +347,59 @@ class InfoCommand (CurveCommand):
         return [block.shape for block in curve.data]
 
 
+class BlockInfoCommand (BlockCommand):
+    """Get selected information about a :class:`hooke.curve.Curve` data block.
+    """
+    def __init__(self, plugin):
+        super(BlockInfoCommand, self).__init__(
+            name='block info', arguments=[
+                Argument(
+                    name='key', count=-1, optional=False,
+                    help='Dot-separted (.) key selection regexp.'),
+                Argument(
+                    name='output',
+                    help="""
+File name for the output (appended).
+""".strip()),
+                ],
+            help=self.__doc__, plugin=plugin)
+
+    def _run(self, hooke, inqueue, outqueue, params):
+        block = self._block(hooke, params)
+        values = {'index': self._block_index(hooke, params)}
+        for key in params['key']:
+            keys = [(0, key.split('.'), block.info)]
+            while len(keys) > 0:
+                index,key_stack,info = keys.pop(0)
+                regexp = re.compile(key_stack[index])
+                matched = False
+                for k,v in info.items():
+                    if regexp.match(k):
+                        matched = True
+                        new_stack = copy.copy(key_stack)
+                        new_stack[index] = k
+                        if index+1 == len(key_stack):
+                            vals = values
+                            for k in new_stack[:-1]:
+                                if k not in vals:
+                                    vals[k] = {}
+                                vals = vals[k]
+                            vals[new_stack[-1]] = v
+                        else:
+                            keys.append((index+1, new_stack, v))
+                if matched == False:
+                    raise ValueError('no match found for %s in %s'
+                                     % (key_stack[index], key))
+        if params['output'] != None:
+            curve = self._curve(hooke, params)
+            with open(params['output'], 'a') as f:
+                yaml.dump({curve.name:{
+                            'path': curve.path,
+                            block.info['name']: values
+                            }}, f)
+        outqueue.put(values)
+
+
 class DeltaCommand (BlockCommand):
     """Get distance information between two points.