pybabel extract: Support multiple keywords

Extend keywords dict to support multiple keywords with same name/arity

Fixes python-babel#1067

babel/messages/extract.py

@@ -324,6 +324,23 @@ def extract_from_file(
                             options, strip_comment_tags))
 
 
+def _tuple_holds_multiple_specs(spec: tuple[int|tuple[int, str], ...]|
+                                      tuple[tuple[int|tuple[int, str], ...]]):
+    """Helper function for extract() which checks whether a given spec tuple
+    contains multiple specs or only one.
+
+    :param spec: a tuple containing a keyword specification or specifications
+    :returns: True if the tuple contains multiple specifications, False otherwise
+    """
+    if spec is None:
+        return False
+    if len(spec) == 1 and not isinstance(spec[0], tuple):
+        return False
+    if len(spec) == 2 and isinstance(spec[1], int):
+        return False
+    return True
+
+
 def _match_messages_against_spec(lineno: int, messages: list[str|None], comments: list[str],
                                  fileobj: _FileObj, spec: tuple[int|tuple[int, str], ...]):
     translatable = []
@@ -463,11 +480,17 @@ def extract(
                 spec = specs[arity]
             except KeyError:
                 continue
-            if spec is None:
-                spec = (1,)
-            result = _match_messages_against_spec(lineno, messages, comments, fileobj, spec)
-            if result is not None:
-                yield result
+            # To maintain backwards compatibility for keyword dicts that only contain
+            # one spec per arity, put any single spec into a tuple.
+            if (spec is None or
+                    (isinstance(spec, tuple) and not _tuple_holds_multiple_specs(spec))):
+                spec = (spec,)
+            for single_spec in spec:
+                if single_spec is None:
+                    single_spec = (1,)
+                result = _match_messages_against_spec(lineno, messages, comments, fileobj, single_spec)
+                if result is not None:
+                    yield result
 
 
 def extract_nothing(

babel/messages/frontend.py

@@ -1166,6 +1166,11 @@ def parse_keywords(strings: Iterable[str] = ()):
     ``(n, 'c')``, meaning that the nth argument should be extracted as context for the
     messages. A ``None`` specification is equivalent to ``(1,)``, extracting the first
     argument.
+
+    A keyword name/number of arguments can map to multiple specifications. If so,
+    the dictionary value will be a tuple containing all the relevant specifications.
+    For backwards compatibility, if there is only one relevant specification,
+    it will be stored directly in the dictionary rather than in a tuple.
     """
     keywords = {}
     for string in strings:
@@ -1176,10 +1181,19 @@ def parse_keywords(strings: Iterable[str] = ()):
             funcname = string
             number = None
             spec = None
-        keywords.setdefault(funcname, {})[number] = spec
+        if funcname in keywords and number in keywords[funcname]:
+            keywords[funcname][number] = keywords[funcname][number] + (spec,)
+        else:
+            keywords.setdefault(funcname, {})[number] = (spec,)
 
-    # For best backwards compatibility, collapse {None: x} into x.
     for k, v in keywords.items():
+        # For best backwards compatibility, if there is only a single spec for a
+        # keyword/number of arguments combination, take the spec out of its
+        # containing tuple and put it directly into the dict.
+        for arity, spec in v.items():
+            if isinstance (spec, tuple) and len(spec) == 1:
+                keywords[k][arity] = spec[0]
+        # For best backwards compatibility, collapse {None: x} into x.
         if set(v) == {None}:
             keywords[k] = v[None]
 

tests/messages/test_extract.py

@@ -493,6 +493,42 @@ def test_different_signatures(self):
         assert messages[0][1] == 'foo'
         assert messages[1][1] == ('hello', 'there')
 
+        def test_multiple_keywords(self):
+            buf = BytesIO(b"""
+        msg1 = _('foo')
+        msg2 = _('bar', 'bars', len(bars))
+        """)
+            keywords = {
+                '_': ((1,), (1, 2))
+            }
+            messages = \
+                list(extract.extract('python', buf, keywords, [], {}))
+
+            assert len(messages) == 3
+            assert messages[0][0] == 'foo'
+            assert messages[1][0] == 'bar'
+            assert messages[2][1] == 'bars'
+
+        def test_multiple_keywords_with_multiple_arities(self):
+            buf = BytesIO(b"""
+msg1 = _('foo')
+msg2 = _('bar', 'bars', len(bars))
+""")
+            keywords = {
+                '_': {
+                    1: (1,),
+                    3: ((1,), (2,)),
+                }
+            }
+            messages = \
+                list(extract.extract('python', buf, keywords, [], {}))
+
+            assert len(messages) == 3
+            assert messages[0][0] == 'foo'
+            assert messages[1][0] == 'bar'
+            assert messages[2][1] == 'bars'
+
+
     def test_empty_string_msgid(self):
         buf = BytesIO(b"""\
 msg = _('')

tests/messages/test_frontend.py

@@ -1496,6 +1496,17 @@ def test_parse_keywords():
     }
 
 
+def test_parse_duplicate_keywords():
+    kw = frontend.parse_keywords(['_', '_:1,2', '_:3,4', 'dgettext:1', 'dgettext:1,2',
+                                 'pgettext:1,2'])
+
+    assert kw == {
+        '_': (None, (1, 2), (3, 4)),
+        'dgettext': ((1,), (1, 2)),
+        'pgettext': (1, 2),
+    }
+
+
 def test_parse_keywords_with_t():
     kw = frontend.parse_keywords(['_:1', '_:2,2t', '_:2c,3,3t'])
 
@@ -1508,6 +1519,17 @@ def test_parse_keywords_with_t():
     }
 
 
+def test_parse_duplicate_keywords_with_t():
+    kw = frontend.parse_keywords(['_:1', '_:1,2', '_:2,3t', '_:3,3t'])
+
+    assert kw == {
+        '_': {
+            None: ((1,), (1,2)),
+            3: ((2,), (3,)),
+        }
+    }
+
+
 def test_extract_messages_with_t():
     content = rb"""
 _("1 arg, arg 1")