Remove registration._get_extension_class() since it's redundant

with registry.class_for_type().

stix2/custom.py

@@ -3,9 +3,10 @@ from collections import OrderedDict
 from .base import _cls_init
 from .properties import EnumProperty
 from .registration import (
-    _get_extension_class, _register_extension, _register_marking,
-    _register_object, _register_observable,
+    _register_extension, _register_marking, _register_object,
+    _register_observable,
 )
+from .registry import class_for_type
 
 
 def _get_properties_dict(properties):
@@ -34,7 +35,7 @@ def _custom_object_builder(cls, type, properties, version, base_class):
             if ext and version != '2.0':
                 if 'extensions' not in self._inner:
                     self._inner['extensions'] = {}
-                self._inner['extensions'][ext] = _get_extension_class(ext, version)()
+                self._inner['extensions'][ext] = class_for_type(ext, version, "extensions")()
 
     _CustomObject.__name__ = cls.__name__
 
@@ -80,7 +81,7 @@ def _custom_observable_builder(cls, type, properties, version, base_class, id_co
             if ext and version != '2.0':
                 if 'extensions' not in self._inner:
                     self._inner['extensions'] = {}
-                self._inner['extensions'][ext] = _get_extension_class(ext, version)()
+                self._inner['extensions'][ext] = class_for_type(ext, version, "extensions")()
 
     _CustomObservable.__name__ = cls.__name__
 

stix2/registration.py

@@ -132,11 +132,6 @@ def _register_observable(new_observable, version=version.DEFAULT_VERSION):
     OBJ_MAP_OBSERVABLE[new_observable._type] = new_observable
 
 
-def _get_extension_class(extension_uuid, version):
-    """Retrieve a registered class Extension"""
-    return registry.STIX2_OBJ_MAPS[version]['extensions'].get(extension_uuid)
-
-
 def _register_extension(
     new_extension, version=version.DEFAULT_VERSION,
 ):