[WIP] Pickle TypeMap

setup.py

@@ -1,10 +1,30 @@
 # -*- coding: utf-8 -*-
 
 from setuptools import setup, find_packages
+from setuptools.command.build_py import build_py
 import re
 
 import versioneer
 
+class BuildPyCommand(build_py):
+    tm_pkl_path = 'typemap.pkl'
+    def run(self):
+        ########## BEGIN: PICKLE TYPE MAP
+        import sys
+        import os
+        import pickle
+
+        sys.path.append('src')
+        import pynwb
+        tm = pynwb.get_type_map()
+
+        pkl_path = os.path.join('src', 'pynwb', self.tm_pkl_path)
+        sys.stdout.write('pickling type map to %s\n' % pkl_path)
+        with open(pkl_path, 'wb') as f:
+            pickle.dump(tm, f, pickle.HIGHEST_PROTOCOL)
+        ########## END: PICKLE TYPE MAP
+        super().run()
+
 with open('README.rst', 'r') as fp:
     readme = fp.read()
 
@@ -33,7 +53,8 @@
     'install_requires': reqs,
     'packages': pkgs,
     'package_dir': {'': 'src'},
-    'package_data': {'pynwb': ["%s/*.yaml" % schema_dir, "%s/*.json" % schema_dir]},
+    #'package_data': {'pynwb': ["%s/*.yaml" % schema_dir, "%s/*.json" % schema_dir]},
+    'package_data': {'pynwb': [BuildPyCommand.tm_pkl_path]},
     'classifiers': [
         "Programming Language :: Python",
         "Programming Language :: Python :: 3.5",
@@ -62,7 +83,8 @@
                 'NWB '
                 'NWB:N '
                 'NeurodataWithoutBorders',
-    'zip_safe': False
+    'zip_safe': False,
+    'cmdclass': {'build_py': BuildPyCommand}
 }
 
 if __name__ == '__main__':

src/pynwb/__init__.py

@@ -1,6 +1,5 @@
-'''This package will contain functions, classes, and objects
-for reading and writing data in NWB format
-'''
+"""This package contains functions, classes, and objects for reading and writing data in NWB format
+"""
 import os.path
 from copy import deepcopy
 from warnings import warn
@@ -14,14 +13,65 @@
 from hdmf.build import BuildManager, TypeMap
 
 
+__rct_kwargs = list()
+
+
+# a function to register a container classes with the global map
+@docval({'name': 'neurodata_type', 'type': str, 'doc': 'the neurodata_type to get the spec for'},
+        {'name': 'namespace', 'type': str, 'doc': 'the name of the namespace'},
+        {"name": "container_cls", "type": type,
+         "doc": "the class to map to the specified neurodata_type", 'default': None},
+        is_method=False)
+def register_class(**kwargs):
+    """Register an NWBContainer class to use for reading and writing a neurodata_type from a specification
+    If container_cls is not specified, returns a decorator for registering an NWBContainer subclass
+    as the class for neurodata_type in namespace.
+    """
+    neurodata_type, namespace, container_cls = getargs('neurodata_type', 'namespace', 'container_cls', kwargs)
+
+    def _dec(cls):
+        __rct_kwargs.append({'data_type': neurodata_type, 'namespace': namespace, 'container_cls': cls})
+        return cls
+    if container_cls is None:
+        return _dec
+    else:
+        _dec(container_cls)
+
+
+__rm_kwargs = list()
+
+
+# a function to register an object mapper for a container class
+@docval({"name": "container_cls", "type": type,
+         "doc": "the Container class for which the given ObjectMapper class gets used for"},
+        {"name": "mapper_cls", "type": type, "doc": "the ObjectMapper class to use to map", 'default': None},
+        is_method=False)
+def register_map(**kwargs):
+    """Register an ObjectMapper to use for a Container class type
+    If mapper_cls is not specified, returns a decorator for registering an ObjectMapper class as the mapper for
+    container_cls. If mapper_cls specified, register the class as the mapper for container_cls.
+    """
+    container_cls, mapper_cls = getargs('container_cls', 'mapper_cls', kwargs)
+
+    def _dec(cls):
+        __rm_kwargs.append({'mapper_cls': cls, 'container_cls': container_cls})
+        return cls
+    if mapper_cls is None:
+        return _dec
+    else:
+        _dec(mapper_cls)
+
+
 CORE_NAMESPACE = 'core'
 __core_ns_file_name = 'nwb.namespace.yaml'
+__typemap_pkl_name = 'typemap.pkl'
 
 
 def __get_resources():
     from pkg_resources import resource_filename
     ret = dict()
     ret['namespace_path'] = os.path.join(resource_filename(__name__, 'nwb-schema/core'), __core_ns_file_name)
+    ret['cached_typemap_path'] = resource_filename(__name__, __typemap_pkl_name)
     return ret
 
 
@@ -30,14 +80,66 @@ def _get_resources():
     return __get_resources()
 
 
-# a global namespace catalog
-global __NS_CATALOG
+# a global type map
 global __TYPE_MAP
 
-from .spec import NWBDatasetSpec, NWBGroupSpec, NWBNamespace  # noqa E402
 
-__NS_CATALOG = NamespaceCatalog(NWBGroupSpec, NWBDatasetSpec, NWBNamespace)
-__TYPE_MAP = TypeMap(__NS_CATALOG)
+@docval({'name': 'namespace_path', 'type': str, 'doc': 'the path to the YAML with the namespace definition'},
+        returns="the namespaces loaded from the given file", rtype=tuple,
+        is_method=False)
+def load_namespaces(**kwargs):
+    """Load namespaces from a file"""
+    namespace_path = getargs('namespace_path', kwargs)
+    return __TYPE_MAP.load_namespaces(namespace_path)
+
+
+def available_namespaces():
+    """Returns all namespaces registered in the namespace catalog"""
+    return __TYPE_MAP.namespace_catalog.namespaces
+
+
+# load the core namespace i.e. base NWB specification
+__resources = __get_resources()
+if os.path.exists(__resources['cached_typemap_path']):
+    import pickle
+    with open(__resources['cached_typemap_path'], 'rb') as f:
+        __TYPE_MAP = pickle.load(f)
+elif os.path.exists(__resources['namespace_path']):
+    from .spec import NWBDatasetSpec, NWBGroupSpec, NWBNamespace  # noqa E402
+    __TYPE_MAP = TypeMap(NamespaceCatalog(NWBGroupSpec, NWBDatasetSpec, NWBNamespace))
+
+    load_namespaces(__resources['namespace_path'])
+
+    # import these so the TypeMap gets populated
+    from . import io as __io  # noqa: F401,E402
+
+    from . import core  # noqa: F401,E402
+    from . import base  # noqa: F401,E402
+    from . import file  # noqa: F401,E402
+    from . import behavior  # noqa: F401,E402
+    from . import device  # noqa: F401,E402
+    from . import ecephys  # noqa: F401,E402
+    from . import epoch  # noqa: F401,E402
+    from . import icephys  # noqa: F401,E402
+    from . import image  # noqa: F401,E402
+    from . import misc  # noqa: F401,E402
+    from . import ogen  # noqa: F401,E402
+    from . import ophys  # noqa: F401,E402
+    from . import retinotopy  # noqa: F401,E402
+
+    for _ in __rct_kwargs:
+        __TYPE_MAP.register_container_type(**_)
+    for _ in __rm_kwargs:
+        __TYPE_MAP.register_map(**_)
+else:
+    raise RuntimeError("Unable to load a TypeMap")
+
+
+NWBContainer = __TYPE_MAP.get_container_cls(CORE_NAMESPACE, 'NWBContainer')
+NWBData = __TYPE_MAP.get_container_cls(CORE_NAMESPACE, 'NWBData')
+TimeSeries = __TYPE_MAP.get_container_cls(CORE_NAMESPACE, 'TimeSeries')
+ProcessingModule = __TYPE_MAP.get_container_cls(CORE_NAMESPACE, 'ProcessingModule')
+NWBFile = __TYPE_MAP.get_container_cls(CORE_NAMESPACE, 'NWBFile')
 
 
 @docval({'name': 'extensions', 'type': (str, TypeMap, list),
@@ -86,72 +188,7 @@ def get_manager(**kwargs):
     return BuildManager(type_map)
 
 
-@docval({'name': 'namespace_path', 'type': str, 'doc': 'the path to the YAML with the namespace definition'},
-        returns="the namespaces loaded from the given file", rtype=tuple,
-        is_method=False)
-def load_namespaces(**kwargs):
-    '''
-    Load namespaces from file
-    '''
-    namespace_path = getargs('namespace_path', kwargs)
-    return __TYPE_MAP.load_namespaces(namespace_path)
-
-
-# load the core namespace i.e. base NWB specification
-__resources = __get_resources()
-if os.path.exists(__resources['namespace_path']):
-    load_namespaces(__resources['namespace_path'])
-
-
-def available_namespaces():
-    """Returns all namespaces registered in the namespace catalog"""
-    return __NS_CATALOG.namespaces
-
-
-# a function to register a container classes with the global map
-@docval({'name': 'neurodata_type', 'type': str, 'doc': 'the neurodata_type to get the spec for'},
-        {'name': 'namespace', 'type': str, 'doc': 'the name of the namespace'},
-        {"name": "container_cls", "type": type, "doc": "the class to map to the specified neurodata_type",
-         'default': None},
-        is_method=False)
-def register_class(**kwargs):
-    """Register an NWBContainer class to use for reading and writing a neurodata_type from a specification
-    If container_cls is not specified, returns a decorator for registering an NWBContainer subclass
-    as the class for neurodata_type in namespace.
-    """
-    neurodata_type, namespace, container_cls = getargs('neurodata_type', 'namespace', 'container_cls', kwargs)
-
-    def _dec(cls):
-        __TYPE_MAP.register_container_type(namespace, neurodata_type, cls)
-        return cls
-    if container_cls is None:
-        return _dec
-    else:
-        _dec(container_cls)
-
-
-# a function to register an object mapper for a container class
-@docval({"name": "container_cls", "type": type,
-         "doc": "the Container class for which the given ObjectMapper class gets used"},
-        {"name": "mapper_cls", "type": type, "doc": "the ObjectMapper class to use to map", 'default': None},
-        is_method=False)
-def register_map(**kwargs):
-    """Register an ObjectMapper to use for a Container class type
-    If mapper_cls is not specified, returns a decorator for registering an ObjectMapper class as the mapper for
-    container_cls. If mapper_cls is specified, register the class as the mapper for container_cls
-    """
-    container_cls, mapper_cls = getargs('container_cls', 'mapper_cls', kwargs)
-
-    def _dec(cls):
-        __TYPE_MAP.register_map(container_cls, cls)
-        return cls
-    if mapper_cls is None:
-        return _dec
-    else:
-        _dec(mapper_cls)
-
-
-@docval({'name': 'neurodata_type', 'type': str, 'doc': 'the neurodata_type to get the NWBContainer class for'},
+@docval({'name': 'neurodata_type', 'type': str, 'doc': 'the neurodata_type to get the NWBConatiner class for'},
         {'name': 'namespace', 'type': str, 'doc': 'the namespace the neurodata_type is defined in'},
         is_method=False)
 def get_class(**kwargs):
@@ -220,23 +257,6 @@ def __init__(self, **kwargs):
         super(NWBHDF5IO, self).__init__(path, manager=manager, mode=mode, file=file_obj, comm=comm)
 
 
-from . import io as __io  # noqa: F401,E402
-from .core import NWBContainer, NWBData  # noqa: F401,E402
-from .base import TimeSeries, ProcessingModule  # noqa: F401,E402
-from .file import NWBFile  # noqa: F401,E402
-
-from . import behavior  # noqa: F401,E402
-from . import device  # noqa: F401,E402
-from . import ecephys  # noqa: F401,E402
-from . import epoch  # noqa: F401,E402
-from . import icephys  # noqa: F401,E402
-from . import image  # noqa: F401,E402
-from . import misc  # noqa: F401,E402
-from . import ogen  # noqa: F401,E402
-from . import ophys  # noqa: F401,E402
-from . import retinotopy  # noqa: F401,E402
-from . import legacy  # noqa: F401,E402
-
 from ._version import get_versions  # noqa: E402
 __version__ = get_versions()['version']
 del get_versions

src/pynwb/legacy/__init__.py

@@ -1,21 +1,15 @@
 
-from hdmf.spec import NamespaceCatalog
 from hdmf.utils import docval, getargs
-from ..spec import NWBDatasetSpec, NWBGroupSpec, NWBNamespace
-from .. import _get_resources, get_type_map, NWBContainer
+from .. import get_type_map, NWBContainer
 
 from .map import ObjectMapperLegacy as ObjectMapper
 from .map import TypeMapLegacy as TypeMap
 
 
-__NS_CATALOG = NamespaceCatalog(NWBGroupSpec, NWBDatasetSpec, NWBNamespace)
-__TYPE_MAP = TypeMap(__NS_CATALOG)
-
 # load the core namespace i.e. base NWB specification
-__resources = _get_resources()
-__TYPE_MAP.load_namespaces(__resources['namespace_path'])
-__TYPE_MAP.merge(get_type_map())
-
+dflt_tm = get_type_map()
+__TYPE_MAP = TypeMap(dflt_tm.namespace_catalog)
+__TYPE_MAP.merge(dflt_tm)
 # Register new ObjectMapper with the new TypeMap:
 __TYPE_MAP.register_map(NWBContainer, ObjectMapper)