Merge pull request #289 from akx/allow-extract-callabledev-2a51c9b95d06

Allow passing a callable to `extract()`

2 files changed, 12 insertions, 2 deletions

diff --git a/babel/messages/extract.py b/babel/messages/extract.py
index 2f8084a..a0608d7 100644
--- a/babel/messages/extract.py
+++ b/babel/messages/extract.py
@@ -212,7 +212,8 @@ def extract(method, fileobj, keywords=DEFAULT_KEYWORDS, comment_tags=(),
     ...     print message
     (3, u'Hello, world!', [], None)
 
-    :param method: a string specifying the extraction method (.e.g. "python");
+    :param method: an extraction method (a callable), or
+                   a string specifying the extraction method (.e.g. "python");
                    if this is a simple name, the extraction function will be
                    looked up by entry point; if it is an explicit reference
                    to a function (of the form ``package.module:funcname`` or
@@ -231,7 +232,9 @@ def extract(method, fileobj, keywords=DEFAULT_KEYWORDS, comment_tags=(),
     :raise ValueError: if the extraction method is not registered
     """
     func = None
-    if ':' in method or '.' in method:
+    if callable(method):
+        func = method
+    elif ':' in method or '.' in method:
         if ':' not in method:
             lastdot = method.rfind('.')
             module, attrname = method[:lastdot], method[lastdot + 1:]
@@ -258,6 +261,7 @@ def extract(method, fileobj, keywords=DEFAULT_KEYWORDS, comment_tags=(),
                 'javascript': extract_javascript
             }
             func = builtin.get(method)
+
     if func is None:
         raise ValueError('Unknown extraction method %r' % method)
 
diff --git a/tests/messages/test_extract.py b/tests/messages/test_extract.py
index ded697f..fa03207 100644
--- a/tests/messages/test_extract.py
+++ b/tests/messages/test_extract.py
@@ -546,3 +546,9 @@ msg = _('')
             assert 'warning: Empty msgid.' in sys.stderr.getvalue()
         finally:
             sys.stderr = stderr
+
+    def test_extract_allows_callable(self):
+        def arbitrary_extractor(fileobj, keywords, comment_tags, options):
+            return [(1, None, (), ())]
+        for x in extract.extract(arbitrary_extractor, BytesIO(b"")):
+            assert x[0] == 1