diff --git a/s2protocol/versions/__init__.py b/s2protocol/versions/__init__.py index aaed865..2f7fb6a 100644 --- a/s2protocol/versions/__init__.py +++ b/s2protocol/versions/__init__.py @@ -1,47 +1,81 @@ - +import logging import os import re -import imp +import importlib import sys - +PROTOCOLS=[] def _import_protocol(base_path, protocol_module_name): """ Import a module from a base path, used to import protocol modules. - This implementation is derived from the __import__ example here: - https://docs.python.org/2/library/imp.html + This implementation is derived from the import_module example here: + https://docs.python.org/3/library/importlib.html#approximating-importlib-import-module """ + requested_protocol = protocol_module_name + if protocol_module_name not in PROTOCOLS: + fallback_module_name = max( + PROTOCOLS, + key=protocol_version + ) + logger = logging.getLogger(__name__) + logger.warning( + "Protocol %s not supported, using %s", + requested_protocol, + fallback_module_name, + ) + protocol_module_name = fallback_module_name.split('.')[0] + + # Try to return the module if it's been loaded already try: return sys.modules[protocol_module_name] except KeyError: pass - # If any of the following calls raises an exception, - # there's a problem we can't handle -- let the caller handle it. - # - fp, pathname, description = imp.find_module(protocol_module_name, [base_path]) - try: - return imp.load_module(protocol_module_name, fp, pathname, description) - finally: - # Since we may exit via an exception, close fp explicitly. - if fp: - fp.close() - + # Otherwise try to load the module + spec = importlib.util.spec_from_file_location( + protocol_module_name, + os.path.join(base_path, protocol_module_name+'.py') + ) + if spec is None or spec.loader is None: + raise ImportError( + f"Unable to load protocol module '{protocol_module_name}'" + ) + module = importlib.util.module_from_spec(spec) + + # If the module does not exist, raise an exception and let the + # caller handle it + if( not os.path.exists(spec.origin) ): + raise ImportError("No module named '{}'".format(protocol_module_name)) + + sys.modules[protocol_module_name] = module + spec.loader.exec_module(module) + return module + +def protocol_version(name): + m = re.search(r'(\d+)', name) + if not m: + raise ValueError(f"Invalid protocol name: {name}") + return int(m.group(1)) def list_all(base_path=None): """ Returns a list of the current protocol version file names in the versions module sorted by name. """ + global PROTOCOLS if base_path is None: base_path = os.path.dirname(__file__) pattern = re.compile('.*protocol[0-9]+.py$') - files = [ f for f in os.listdir(base_path) \ - if pattern.match(f) ] - files.sort() - return files + PROTOCOLS = [ f for f in os.listdir(base_path) \ + if pattern.match(f) ] + + """ + sort string containing integers is problematic. + """ + PROTOCOLS.sort(key=protocol_version) + return PROTOCOLS def latest():