Repository: pystorm/pystorm

Issue:
Have Tuple.values be a namedtuple so fields can be accessed by name
_From @dan-blanchard on April 15, 2015 13:57_

This was brought up as part of our discussion of the rejected #120.  What we want to do is:

- [x]  Submit a PR to Storm that serializes  [`TopologyContext.componentToStreamToFields`](https://github.com/apache/storm/blob/master/storm-core/src/jvm/backtype/storm/task/TopologyContext.java#L61) and sends that along as part of Multi-Lang handshake.
- [x]  Add a `_stream_fields` dictionary attribute to the  `Component` class that maps from streams to `namedtuple` types representing the names of fields/values in the tuple.  This should get created at handshake time based on the contents of `componentToStreamToFields`.
- [x]  Modify `Component.read_tuple()` to set `Tuple.values` to be a `namedtuple` of the appropriate type for the current stream (by looking it up in `Component._stream_fields`.

This will allow users to get values out of their tuples by accessing values directly by name (`word = tup.values.word`), or by unpacking (`word, count = tup.values`).

_Copied from original issue: Parsely/streamparse#127_

Discussion:

Patch:
diff --git a/pystorm/component.py b/pystorm/component.py
index 23d82e8..80c5abf 100644
--- a/pystorm/component.py
+++ b/pystorm/component.py
@@ -3,14 +3,17 @@ from __future__ import absolute_import, print_function, unicode_literals
 
 import logging
 import os
+import re
 import signal
 import sys
-from collections import deque, namedtuple
+from collections import defaultdict, deque, namedtuple
 from logging.handlers import RotatingFileHandler
 from os.path import join
 from threading import RLock
 from traceback import format_exc
 
+from six import iteritems
+
 from .exceptions import StormWentAwayError
 from .serializers.msgpack_serializer import MsgpackSerializer
 from .serializers.json_serializer import JSONSerializer
@@ -36,6 +39,9 @@ _PYTHON_LOG_LEVELS = {'critical': logging.CRITICAL,
                       'debug': logging.DEBUG,
                       'trace': logging.DEBUG}
 _SERIALIZERS = {"json": JSONSerializer, "msgpack": MsgpackSerializer}
+# Convert names to valid Python identifiers by replacing non-word characters
+# whitespace and leading digits with underscores.
+_IDENTIFIER_RE = re.compile(r'\W|^(?=\d)')
 
 
 log = logging.getLogger(__name__)
@@ -121,7 +127,7 @@ Tuple = namedtuple('Tuple', 'id component stream task values')
 :ivar task: the task the Tuple was generated from.
 :type task: int
 :ivar values: the payload of the Tuple where data is stored.
-:type values: list
+:type values: tuple (or namedtuple for Storm 0.10.0+)
 """
 
 
@@ -177,6 +183,7 @@ class Component(object):
         self.context = None
         self.pid = os.getpid()
         self.logger = None
+        self._source_tuple_types = defaultdict(dict)
         # pending commands/Tuples we read while trying to read task IDs
         self._pending_commands = deque()
         # pending task IDs we read while trying to read commands/Tuples
@@ -207,6 +214,15 @@ class Component(object):
         self.topology_name = storm_conf.get('topology.name', '')
         self.task_id = context.get('taskid', '')
         self.component_name = context.get('componentid')
+        # source->stream->fields requires Storm 0.10.0 or later
+        source_stream_fields = context.get('source->stream->fields', {})
+        for source, stream_fields in iteritems(source_stream_fields):
+            for stream, fields in iteritems(stream_fields):
+                type_name = (_IDENTIFIER_RE.sub('_', source.title()) +
+                             _IDENTIFIER_RE.sub('_', stream.title()) +
+                             'Tuple')
+                self._source_tuple_types[source][stream] = namedtuple(type_name,
+                                                                      fields)
         # If using Storm before 0.10.0 componentid is not available
         if self.component_name is None:
             self.component_name = context.get('task->component', {})\
@@ -280,8 +296,12 @@ class Component(object):
 
     def read_tuple(self):
         cmd = self.read_command()
-        return Tuple(cmd['id'], cmd['comp'], cmd['stream'], cmd['task'],
-                     cmd['tuple'])
+        source = cmd['comp']
+        stream = cmd['stream']
+        values = cmd['tuple']
+        val_type = self._source_tuple_types[source].get(stream)
+        return Tuple(cmd['id'], source, stream, cmd['task'],
+                     tuple(values) if val_type is None else val_type(*values))
 
     def read_handshake(self):
         """Read and process an initial handshake message from Storm."""


Test patch:
diff --git a/test/pystorm/test_bolt.py b/test/pystorm/test_bolt.py
index f9586c3..80cf5c8 100644
--- a/test/pystorm/test_bolt.py
+++ b/test/pystorm/test_bolt.py
@@ -34,7 +34,7 @@ class BoltTests(unittest.TestCase):
         tup_json = "{}\nend\n".format(json.dumps(self.tup_dict)).encode('utf-8')
         self.tup = Tuple(self.tup_dict['id'], self.tup_dict['comp'],
                          self.tup_dict['stream'], self.tup_dict['task'],
-                         self.tup_dict['tuple'],)
+                         tuple(self.tup_dict['tuple']),)
         self.bolt = Bolt(input_stream=BytesIO(tup_json),
                          output_stream=BytesIO())
         self.bolt.initialize({}, {})
@@ -190,7 +190,7 @@ class BoltTests(unittest.TestCase):
     def test_heartbeat_response(self, send_message_mock, read_tuple_mock):
         # Make sure we send sync for heartbeats
         read_tuple_mock.return_value = Tuple(id='foo', task=-1,
-                                             stream='__heartbeat', values=[],
+                                             stream='__heartbeat', values=(),
                                              component='__system')
         self.bolt._run()
         send_message_mock.assert_called_with(self.bolt, {'command': 'sync'})
@@ -201,7 +201,7 @@ class BoltTests(unittest.TestCase):
         # Make sure we send sync for heartbeats
         read_tuple_mock.return_value = Tuple(id=None, task=-1,
                                              component='__system',
-                                             stream='__tick', values=[50])
+                                             stream='__tick', values=(50,))
         self.bolt._run()
         process_tick_mock.assert_called_with(self.bolt,
                                              read_tuple_mock.return_value)
@@ -239,8 +239,8 @@ class BatchingBoltTests(unittest.TestCase):
         tups_json = '\nend\n'.join([json.dumps(tup_dict) for tup_dict in
                                     self.tup_dicts] + [''])
         self.tups = [Tuple(tup_dict['id'], tup_dict['comp'], tup_dict['stream'],
-                           tup_dict['task'], tup_dict['tuple']) for tup_dict in
-                     self.tup_dicts]
+                           tup_dict['task'], tuple(tup_dict['tuple']))
+                     for tup_dict in self.tup_dicts]
         self.nontick_tups = [tup for tup in self.tups if tup.stream != '__tick']
         self.bolt = BatchingBolt(input_stream=BytesIO(tups_json.encode('utf-8')),
                                  output_stream=BytesIO())
@@ -364,7 +364,7 @@ class BatchingBoltTests(unittest.TestCase):
     def test_heartbeat_response(self, send_message_mock, read_tuple_mock):
         # Make sure we send sync for heartbeats
         read_tuple_mock.return_value = Tuple(id='foo', task=-1,
-                                             stream='__heartbeat', values=[],
+                                             stream='__heartbeat', values=(),
                                              component='__system')
         self.bolt._run()
         send_message_mock.assert_called_with(self.bolt, {'command': 'sync'})
@@ -375,7 +375,7 @@ class BatchingBoltTests(unittest.TestCase):
         # Make sure we send sync for heartbeats
         read_tuple_mock.return_value = Tuple(id=None, task=-1,
                                              component='__system',
-                                             stream='__tick', values=[50])
+                                             stream='__tick', values=(50,))
         self.bolt._run()
         process_tick_mock.assert_called_with(self.bolt,
                                              read_tuple_mock.return_value)
diff --git a/test/pystorm/test_component.py b/test/pystorm/test_component.py
index f7228f4..c508ae1 100644
--- a/test/pystorm/test_component.py
+++ b/test/pystorm/test_component.py
@@ -7,6 +7,7 @@ from __future__ import absolute_import, print_function, unicode_literals
 import logging
 import os
 import unittest
+from collections import namedtuple
 from io import BytesIO
 
 import simplejson as json
@@ -24,49 +25,48 @@ log = logging.getLogger(__name__)
 
 
 class ComponentTests(unittest.TestCase):
-
-    def test_read_handshake(self):
-        handshake_dict = {
-            "conf": {
-                "topology.message.timeout.secs": 3,
-                "topology.tick.tuple.freq.secs": 1,
-                "topology.debug": True
-            },
-            "pidDir": ".",
-            "context": {
-                "task->component": {
-                    "1": "example-spout",
-                    "2": "__acker",
-                    "3": "example-bolt1",
-                    "4": "example-bolt2"
-                },
-                "taskid": 3,
-                # Everything below this line is only available in Storm 0.10.0+
-                "componentid": "example-bolt1",
-                "stream->target->grouping": {
-                    "default": {
-                        "example-bolt2": {
-                            "type": "SHUFFLE"
-                        }
-                    }
-                },
-                "streams": ["default"],
-                "stream->outputfields": {"default": ["word"]},
-                "source->stream->grouping": {
-                    "example-spout": {
-                        "default": {
-                            "type": "FIELDS",
-                            "fields": ["word"]
-                        }
-                    }
-                },
-                "source->stream->fields": {
-                    "example-spout": {
-                        "default": ["word"]
-                    }
+    conf = {"topology.message.timeout.secs": 3,
+            "topology.tick.tuple.freq.secs": 1,
+            "topology.debug": True,
+            "topology.name": "foo"}
+    context = {
+        "task->component": {
+            "1": "example-spout",
+            "2": "__acker",
+            "3": "example-bolt1",
+            "4": "example-bolt2"
+        },
+        "taskid": 3,
+        # Everything below this line is only available in Storm 0.11.0+
+        "componentid": "example-bolt1",
+        "stream->target->grouping": {
+            "default": {
+                "example-bolt2": {
+                    "type": "SHUFFLE"
+                }
+            }
+        },
+        "streams": ["default"],
+        "stream->outputfields": {"default": ["word"]},
+        "source->stream->grouping": {
+            "example-spout": {
+                "default": {
+                    "type": "FIELDS",
+                    "fields": ["word"]
                 }
             }
+        },
+        "source->stream->fields": {
+            "example-spout": {
+                "default": ["sentence", "word", "number"]
+            }
         }
+    }
+
+    def test_read_handshake(self):
+        handshake_dict = {"conf": self.conf,
+                          "pidDir": ".",
+                          "context": self.context}
         pid_dir = handshake_dict['pidDir']
         expected_conf = handshake_dict['conf']
         expected_context = handshake_dict['context']
@@ -84,52 +84,18 @@ class ComponentTests(unittest.TestCase):
                          component.serializer.output_stream.buffer.getvalue())
 
     def test_setup_component(self):
-        conf = {"topology.message.timeout.secs": 3,
-                "topology.tick.tuple.freq.secs": 1,
-                "topology.debug": True,
-                "topology.name": "foo"}
-        context = {
-            "task->component": {
-                "1": "example-spout",
-                "2": "__acker",
-                "3": "example-bolt1",
-                "4": "example-bolt2"
-            },
-            "taskid": 3,
-            # Everything below this line is only available in Storm 0.11.0+
-            "componentid": "example-bolt1",
-            "stream->target->grouping": {
-                "default": {
-                    "example-bolt2": {
-                        "type": "SHUFFLE"
-                    }
-                }
-            },
-            "streams": ["default"],
-            "stream->outputfields": {"default": ["word"]},
-            "source->stream->grouping": {
-                "example-spout": {
-                    "default": {
-                        "type": "FIELDS",
-                        "fields": ["word"]
-                    }
-                }
-            },
-            "source->stream->fields": {
-                "example-spout": {
-                    "default": ["word"]
-                }
-            }
-        }
+        conf = self.conf
         component = Component(input_stream=BytesIO(),
                               output_stream=BytesIO())
-        component._setup_component(conf, context)
+        component._setup_component(conf, self.context)
+        self.assertEqual(component._source_tuple_types['example-spout']['default'].__name__,
+                         'Example_SpoutDefaultTuple')
         self.assertEqual(component.topology_name, conf['topology.name'])
-        self.assertEqual(component.task_id, context['taskid'])
+        self.assertEqual(component.task_id, self.context['taskid'])
         self.assertEqual(component.component_name,
-                         context['task->component'][str(context['taskid'])])
+                         self.context['task->component'][str(self.context['taskid'])])
         self.assertEqual(component.storm_conf, conf)
-        self.assertEqual(component.context, context)
+        self.assertEqual(component.context, self.context)
 
     def test_read_message(self):
         inputs = [# Task IDs
@@ -259,7 +225,7 @@ class ComponentTests(unittest.TestCase):
         for msg in inputs[::2]:
             output = json.loads(msg)
             output['component'] = output['comp']
-            output['values'] = output['tuple']
+            output['values'] = tuple(output['tuple'])
             del output['comp']
             del output['tuple']
             outputs.append(Tuple(**output))
@@ -272,6 +238,38 @@ class ComponentTests(unittest.TestCase):
             tup = component.read_tuple()
             self.assertEqual(output, tup)
 
+    def test_read_tuple_named_fields(self):
+        # This is only valid for bolts, so we only need to test with task IDs
+        # and Tuples
+        inputs = [('{ "id": "-6955786537413359385", "comp": "example-spout", '
+                   '"stream": "default", "task": 9, "tuple": ["snow white and '
+                   'the seven dwarfs", "field2", 3]}\n'), 'end\n']
+
+        component = Component(input_stream=BytesIO(''.join(inputs).encode('utf-8')),
+                              output_stream=BytesIO())
+        component._setup_component(self.conf, self.context)
+
+        Example_SpoutDefaultTuple = namedtuple('Example_SpoutDefaultTuple',
+                                       field_names=['sentence', 'word',
+                                                    'number'])
+
+        outputs = []
+        for msg in inputs[::2]:
+            output = json.loads(msg)
+            output['component'] = output['comp']
+            output['values'] = Example_SpoutDefaultTuple(*output['tuple'])
+            del output['comp']
+            del output['tuple']
+            outputs.append(Tuple(**output))
+
+        for output in outputs:
+            log.info('Checking Tuple for %r', output)
+            tup = component.read_tuple()
+            self.assertEqual(output.values.sentence, tup.values.sentence)
+            self.assertEqual(output.values.word, tup.values.word)
+            self.assertEqual(output.values.number, tup.values.number)
+            self.assertEqual(output, tup)
+
     def test_send_message(self):
         component = Component(input_stream=BytesIO(), output_stream=BytesIO())
         inputs = [{"command": "emit", "id": 4, "stream": "", "task": 9,