Commits

Gregory Petukhov committed 743a5b6

Enchance Data object to support multiple items

  • Participants
  • Parent commits 93856b8

Comments (0)

Files changed (3)

File grab/spider/base.py

 
     def find_data_handler(self, data):
         try:
-            handler = getattr(self, 'data_%s' % data.name)
+            handler = getattr(self, 'data_%s' % data.data_name)
         except AttributeError:
-            raise NoDataHandler('No handler defined for Data %s' % data.name)
+            raise NoDataHandler('No handler defined for Data %s' % data.data_name)
         else:
             return handler
 
             try:
                 handler(**result.storage)
             except Exception, ex:
-                self.process_handler_error('data_%s' % result.name, ex, task)
+                self.process_handler_error('data_%s' % result.data_name, ex, task)
         elif result is None:
             pass
         else:

File grab/spider/data.py

+NULL = object()
+
 class Data(object):
     """
     Task handlers should return instances of that class.
     """
 
     def __init__(self, data_name, **kwargs):
-        self.name = data_name
+        self.data_name = data_name
         self.storage = kwargs
+
+    def __getitem__(self, key):
+        return self.storage[key]
+
+    def get(self, key, default=NULL):
+        try:
+            return self.storage[key]
+        except KeyError:
+            if default is NULL:
+                raise
+            else:
+                return default

File test/spider_data.py

         bot.add_task(Task('page', url=SERVER.BASE_URL))
         bot.run()
         self.assertEqual(bot.data_processed, [1, 2, {'bar': 'gaz'}])
+
+    def test_data_object_dict_interface(self):
+        data = Data('person', person={'age': 22})
+        self.assertRaises(KeyError, lambda: data['name'])
+        self.assertEqual(data['person'], {'age': 22})
+
+    def test_data_object_get_method(self):
+        data = Data('person', person={'age': 22})
+        self.assertRaises(KeyError, lambda: data.get('name'))
+        self.assertEqual('foo', data.get('name', 'foo'))
+        self.assertEqual({'age': 22}, data.get('person'))