Fix workbench wrapped classes for `parse()`.

The wrapped classes need to be in the OBJ_MAP mapping, not just the
workbench.py globals.
stix2.0
Chris Lenk 2018-03-23 13:13:41 -04:00
parent b9bbd03481
commit 98cc86eef6
1 changed files with 19 additions and 10 deletions

View File

@ -21,6 +21,7 @@
""" """
import stix2
from . import AttackPattern as _AttackPattern from . import AttackPattern as _AttackPattern
from . import Campaign as _Campaign from . import Campaign as _Campaign
from . import CourseOfAction as _CourseOfAction from . import CourseOfAction as _CourseOfAction
@ -121,20 +122,28 @@ def _constructor_wrapper(obj_type):
@staticmethod @staticmethod
def new_constructor(cls, *args, **kwargs): def new_constructor(cls, *args, **kwargs):
return _environ.create(wrapped_type, *args, **kwargs) x = _environ.create(wrapped_type, *args, **kwargs)
return x
return new_constructor return new_constructor
# Create wrapper classes whose constructors call the implicit environment's create() def _setup_workbench():
for obj_type in STIX_OBJS: # Create wrapper classes whose constructors call the implicit environment's create()
new_class_dict = { for obj_type in STIX_OBJS:
'__new__': _constructor_wrapper(obj_type), new_class_dict = {
'__doc__': 'Workbench wrapper around the `{0} <stix2.v20.sdo.html#stix2.v20.sdo.{0}>`__. object. {1}'.format(obj_type.__name__, STIX_OBJ_DOCS) '__new__': _constructor_wrapper(obj_type),
} '__doc__': 'Workbench wrapper around the `{0} <stix2.v20.sdo.html#stix2.v20.sdo.{0}>`__. object. {1}'.format(obj_type.__name__, STIX_OBJ_DOCS)
new_class = type(obj_type.__name__, (), new_class_dict) }
new_class = type(obj_type.__name__, (), new_class_dict)
globals()[obj_type.__name__] = new_class # Add our new class to this module's globals and to the library-wide mapping.
new_class = None # This allows parse() to use the wrapped classes.
globals()[obj_type.__name__] = new_class
stix2.OBJ_MAP[obj_type._type] = new_class
new_class = None
_setup_workbench()
# Functions to get all objects of a specific type # Functions to get all objects of a specific type