Add reading of setuptools metadata to find smart_open transport / compressor extensions

Author: arthurlm

smart_open/compression.py

@@ -7,6 +7,8 @@
 #
 """Implements the compression layer of the ``smart_open`` library."""
 import logging
+import importlib
+import importlib.metadata
 import os.path
 
 logger = logging.getLogger(__name__)
@@ -145,3 +147,16 @@ def compression_wrapper(file_obj, mode, compression):
 #
 register_compressor('.bz2', _handle_bz2)
 register_compressor('.gz', _handle_gzip)
+
+
+def _register_compressor_entry_point(ep):
+    try:
+        assert len(ep.name) > 0, "At least one char is required for ep.name"
+        extension = ".{}".format(ep.name)
+        register_compressor(extension, ep.load())
+    except Exception:
+        logger.warning("Fail to load smart_open compressor extension: %s (target: %s)", ep.name, ep.value)
+
+
+for ep in importlib.metadata.entry_points().select(group='smart_open_compressor'):
+    _register_compressor_entry_point(ep)

smart_open/tests/fixtures/compressor.py

@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+"""Some no-op compressor"""
+
+
+def handle_foo():
+    ...
+
+
+def handle_bar():
+    ...

smart_open/tests/test_compression.py

@@ -0,0 +1,53 @@
+# -*- coding: utf-8 -*-
+
+from importlib.metadata import EntryPoint
+import pytest
+
+from smart_open.compression import _COMPRESSOR_REGISTRY, _register_compressor_entry_point
+
+
+def unregister_compressor(ext):
+    if ext in _COMPRESSOR_REGISTRY:
+        del _COMPRESSOR_REGISTRY[ext]
+
+
+@pytest.fixture(autouse=True)
+def cleanup_compressor():
+    unregister_compressor(".foo")
+    unregister_compressor(".bar")
+
+
+def test_register_valid_entry_point():
+    assert ".foo" not in _COMPRESSOR_REGISTRY
+    assert ".bar" not in _COMPRESSOR_REGISTRY
+    _register_compressor_entry_point(EntryPoint(
+        "foo",
+        "smart_open.tests.fixtures.compressor:handle_bar",
+        "smart_open_compressor",
+    ))
+    _register_compressor_entry_point(EntryPoint(
+        "bar",
+        "smart_open.tests.fixtures.compressor:handle_bar",
+        "smart_open_compressor",
+    ))
+    assert ".foo" in _COMPRESSOR_REGISTRY
+    assert ".bar" in _COMPRESSOR_REGISTRY
+
+
+def test_register_invalid_entry_point_name_do_not_crash():
+    _register_compressor_entry_point(EntryPoint(
+        "",
+        "smart_open.tests.fixtures.compressor:handle_foo",
+        "smart_open_compressor",
+    ))
+    assert "" not in _COMPRESSOR_REGISTRY
+    assert "." not in _COMPRESSOR_REGISTRY
+
+
+def test_register_invalid_entry_point_value_do_not_crash():
+    _register_compressor_entry_point(EntryPoint(
+        "foo",
+        "smart_open.tests.fixtures.compressor:handle_invalid",
+        "smart_open_compressor",
+    ))
+    assert ".foo" not in _COMPRESSOR_REGISTRY

smart_open/tests/test_transport.py

@@ -1,16 +1,44 @@
 # -*- coding: utf-8 -*-
+from importlib.metadata import EntryPoint
 import pytest
 import unittest
 
-from smart_open.transport import register_transport, get_transport
+from smart_open.transport import (
+    register_transport, get_transport, _REGISTRY, _ERRORS, _register_transport_entry_point
+)
+
+
+def unregister_transport(x):
+    if x in _REGISTRY:
+        del _REGISTRY[x]
+    if x in _ERRORS:
+        del _ERRORS[x]
+
+
+def assert_transport_not_registered(scheme):
+    with pytest.raises(NotImplementedError):
+        get_transport(scheme)
+
+
+def assert_transport_registered(scheme):
+    transport = get_transport(scheme)
+    assert transport.SCHEME == scheme
 
 
 class TransportTest(unittest.TestCase):
+    def tearDown(self):
+        unregister_transport("foo")
+        unregister_transport("missing")
 
     def test_registry_requires_declared_schemes(self):
         with pytest.raises(ValueError):
             register_transport('smart_open.tests.fixtures.no_schemes_transport')
 
+    def test_registry_valid_transport(self):
+        assert_transport_not_registered("foo")
+        register_transport('smart_open.tests.fixtures.good_transport')
+        assert_transport_registered("foo")
+
     def test_registry_errors_on_double_register_scheme(self):
         register_transport('smart_open.tests.fixtures.good_transport')
         with pytest.raises(AssertionError):
@@ -20,3 +48,29 @@ def test_registry_errors_get_transport_for_module_with_missing_deps(self):
         register_transport('smart_open.tests.fixtures.missing_deps_transport')
         with pytest.raises(ImportError):
             get_transport("missing")
+
+    def test_register_entry_point_valid(self):
+        assert_transport_not_registered("foo")
+        _register_transport_entry_point(EntryPoint(
+            "foo",
+            "smart_open.tests.fixtures.good_transport",
+            "smart_open_transport",
+        ))
+        assert_transport_registered("foo")
+
+    def test_register_entry_point_catch_bad_data(self):
+        _register_transport_entry_point(EntryPoint(
+            "invalid",
+            "smart_open.some_totaly_invalid_module",
+            "smart_open_transport",
+        ))
+
+    def test_register_entry_point_for_module_with_missing_deps(self):
+        assert_transport_not_registered("missing")
+        _register_transport_entry_point(EntryPoint(
+            "missing",
+            "smart_open.tests.fixtures.missing_deps_transport",
+            "smart_open_transport",
+        ))
+        with pytest.raises(ImportError):
+            get_transport("missing")

smart_open/transport.py

@@ -11,6 +11,7 @@
 
 """
 import importlib
+import importlib.metadata
 import logging
 
 import smart_open.local_file
@@ -102,5 +103,16 @@ def get_transport(scheme):
 register_transport('smart_open.ssh')
 register_transport('smart_open.webhdfs')
 
+
+def _register_transport_entry_point(ep):
+    try:
+        register_transport(ep.value)
+    except Exception:
+        logger.warning("Fail to load smart_open transport extension: %s (target: %s)", ep.name, ep.value)
+
+
+for ep in importlib.metadata.entry_points().select(group='smart_open_transport'):
+    _register_transport_entry_point(ep)
+
 SUPPORTED_SCHEMES = tuple(sorted(_REGISTRY.keys()))
 """The transport schemes that the local installation of ``smart_open`` supports."""