def test_bad_generator(self):
"""Test that register_generator raises an error if the generator is not an instance of CustomClassGenerator."""
class NotACustomClassGenerator:
pass
type_map = TypeMap()
msg = 'Generator <.*> must be a subclass of CustomClassGenerator.'
with self.assertRaisesRegex(ValueError, msg):
type_map.register_generator(NotACustomClassGenerator)
def test_register_generator(self):
"""Test TypeMap.register_generator and ClassGenerator.register_generator."""
class MyClassGenerator(CustomClassGenerator):
@classmethod
def apply_generator_to_field(cls, field_spec, bases, type_map):
return True
@classmethod
def process_field_spec(cls, classdict, docval_args, parent_cls,
attr_name, not_inherited_fields, type_map,
spec):
# append attr_name to classdict['__custom_fields__'] list
classdict.setdefault('process_field_spec',
list()).append(attr_name)
@classmethod
def post_process(cls, classdict, bases, docval_args, spec):
classdict['post_process'] = True
spec = GroupSpec(doc='A test group specification with a data type',
data_type_def='Baz',
attributes=[
AttributeSpec(name='attr1',
doc='a string attribute',
dtype='text')
])
spec_catalog = SpecCatalog()
spec_catalog.register_spec(spec, 'test.yaml')
namespace = SpecNamespace(doc='a test namespace',
name=CORE_NAMESPACE,
schema=[{
'source': 'test.yaml'
}],
version='0.1.0',
catalog=spec_catalog)
namespace_catalog = NamespaceCatalog()
namespace_catalog.add_namespace(CORE_NAMESPACE, namespace)
type_map = TypeMap(namespace_catalog)
type_map.register_generator(MyClassGenerator)
cls = type_map.get_dt_container_cls('Baz', CORE_NAMESPACE)
self.assertEqual(cls.process_field_spec, ['attr1'])
self.assertTrue(cls.post_process)
