Implement YAML constructor and modify attribute handling

pyworker/job.py

@@ -14,6 +14,12 @@ def __new__(meta, name, bases, class_dict):
         _register_class(cls)
         return cls
 
+# Add a YAML constructor to ignore Ruby-specific tags (required once at module load time)
+def no_ruby_objects(loader, tag_suffix, node):
+    # Construct mapping normally, ignoring Ruby-specific tags
+    return loader.construct_mapping(node)
+
+yaml.SafeLoader.add_multi_constructor("!ruby/object:", no_ruby_objects)
 
 class Job(object):
     """docstring for Job"""
@@ -52,19 +58,6 @@ def extract_class_name(line):
             else:
                 return None
 
-        def extract_attributes(lines):
-            attributes = []
-            collect = False
-            for line in lines:
-                if line.startswith('  raw_attributes:'):
-                    collect = True
-                elif not line.startswith('    '):
-                    if collect:
-                        break
-                elif collect:
-                    attributes.append(line)
-            return attributes
-
         def extract_extra_fields(extra_fields, extra_field_values):
             if extra_fields is None or extra_field_values is None:
                 return None
@@ -91,18 +84,18 @@ def extract_extra_fields(extra_fields, extra_field_values):
                 abstract=True, extra_fields=extra_fields_dict,
                 reporter=reporter)
 
-        attributes = extract_attributes(handler[2:])
+        attributes = handler[3:]
         logger.debug("Found attributes: %s" % str(attributes))
 
-        stripped = '\n'.join(['object:', '  attributes:'] + attributes)
-        payload = yaml.load(stripped, Loader=yaml.FullLoader)
+        stripped = '\n'.join(['object:', '  raw_attributes:'] + attributes)
+        payload = yaml.load(stripped, Loader=yaml.SafeLoader)
         logger.debug("payload object: %s" % str(payload))
 
         return target_class(class_name=class_name, logger=logger,
             job_id=job_id, attempts=attempts,
             run_at=run_at, queue=queue, database=database,
             max_attempts=max_attempts,
-            attributes=payload['object']['attributes'],
+            attributes=payload['object']['raw_attributes'],
             abstract=False, extra_fields=extra_fields_dict,
             reporter=reporter)
 

tests/fixtures/handler_registered.yaml

@@ -8,6 +8,16 @@ object: !ruby/object:RegisteredJob
       multiline
     total_articles: 1000
     is_blind: true
+    proper_multiline: |-
+      line one
+      line two
+
+      line four
+    collapsed_first_newline: 'line one
+      abc'
+    extra_new_line: 'line one
+
+      line two'
   attributes: !ruby/object:ActiveRecord::AttributeSet
     attributes: !ruby/object:ActiveRecord::LazyAttributeHash
       types: {}

tests/test_job.py

@@ -146,7 +146,10 @@ def test_from_row_when_registered_class_returns_job_instance_with_attributes(sel
             'title': 'review title',
             'description': 'review description\nmultiline\n',
             'total_articles': 1000,
-            'is_blind': True
+            'is_blind': True,
+            'collapsed_first_newline': 'line one abc',
+            'extra_new_line': 'line one\nline two',
+            'proper_multiline': 'line one\nline two\n\nline four'
         })
         self.assertIsNone(job.reporter)