bpo-32222: Fix pygettext skipping docstrings for funcs with arg typehints (GH-4745)

1 files changed, 88 insertions, 0 deletions

diff --git a/Lib/test/test_tools/test_i18n.py b/Lib/test/test_tools/test_i18n.py
index 5c28bae..3c8be27 100644
--- a/Lib/test/test_tools/test_i18n.py
+++ b/Lib/test/test_tools/test_i18n.py
@@ -3,6 +3,7 @@
 import os
 import sys
 import unittest
+import textwrap
 
 from test.support.script_helper import assert_python_ok
 from test.test_tools import skip_if_missing, toolsdir
@@ -28,6 +29,41 @@ class Test_pygettext(unittest.TestCase):
             headers[key] = val.strip()
         return headers
 
+    def get_msgids(self, data):
+        """ utility: return all msgids in .po file as a list of strings """
+        msgids = []
+        reading_msgid = False
+        cur_msgid = []
+        for line in data.split('\n'):
+            if reading_msgid:
+                if line.startswith('"'):
+                    cur_msgid.append(line.strip('"'))
+                else:
+                    msgids.append('\n'.join(cur_msgid))
+                    cur_msgid = []
+                    reading_msgid = False
+                    continue
+            if line.startswith('msgid '):
+                line = line[len('msgid '):]
+                cur_msgid.append(line.strip('"'))
+                reading_msgid = True
+        else:
+            if reading_msgid:
+                msgids.append('\n'.join(cur_msgid))
+
+        return msgids
+
+    def extract_docstrings_from_str(self, module_content):
+        """ utility: return all msgids extracted from module_content """
+        filename = 'test_docstrings.py'
+        with temp_cwd(None) as cwd:
+            with open(filename, 'w') as fp:
+                fp.write(module_content)
+            assert_python_ok(self.script, '-D', filename)
+            with open('messages.pot') as fp:
+                data = fp.read()
+        return self.get_msgids(data)
+
     def test_header(self):
         """Make sure the required fields are in the header, according to:
            http://www.gnu.org/software/gettext/manual/gettext.html#Header-Entry
@@ -72,3 +108,55 @@ class Test_pygettext(unittest.TestCase):
 
             # This will raise if the date format does not exactly match.
             datetime.strptime(creationDate, '%Y-%m-%d %H:%M%z')
+
+    def test_funcdocstring_annotated_args(self):
+        """ Test docstrings for functions with annotated args """
+        msgids = self.extract_docstrings_from_str(textwrap.dedent('''\
+        def foo(bar: str):
+            """doc"""
+        '''))
+        self.assertIn('doc', msgids)
+
+    def test_funcdocstring_annotated_return(self):
+        """ Test docstrings for functions with annotated return type """
+        msgids = self.extract_docstrings_from_str(textwrap.dedent('''\
+        def foo(bar) -> str:
+            """doc"""
+        '''))
+        self.assertIn('doc', msgids)
+
+    def test_funcdocstring_defvalue_args(self):
+        """ Test docstring for functions with default arg values """
+        msgids = self.extract_docstrings_from_str(textwrap.dedent('''\
+        def foo(bar=()):
+            """doc"""
+        '''))
+        self.assertIn('doc', msgids)
+
+    def test_funcdocstring_multiple_funcs(self):
+        """ Test docstring extraction for multiple functions combining
+        annotated args, annotated return types and default arg values
+        """
+        msgids = self.extract_docstrings_from_str(textwrap.dedent('''\
+        def foo1(bar: tuple=()) -> str:
+            """doc1"""
+
+        def foo2(bar: List[1:2]) -> (lambda x: x):
+            """doc2"""
+
+        def foo3(bar: 'func'=lambda x: x) -> {1: 2}:
+            """doc3"""
+        '''))
+        self.assertIn('doc1', msgids)
+        self.assertIn('doc2', msgids)
+        self.assertIn('doc3', msgids)
+
+    def test_classdocstring_early_colon(self):
+        """ Test docstring extraction for a class with colons occuring within
+        the parentheses.
+        """
+        msgids = self.extract_docstrings_from_str(textwrap.dedent('''\
+        class D(L[1:2], F({1: 2}), metaclass=M(lambda x: x)):
+            """doc"""
+        '''))
+        self.assertIn('doc', msgids)