Fix handling of docstrings with tokenization errors

Author: gareth-cross

mypy/stubdoc.py

@@ -347,14 +347,22 @@ def infer_sig_from_docstring(docstr: str | None, name: str) -> list[FunctionSig]
         return None
 
     state = DocStringParser(name)
-    # Return all found signatures, even if there is a parse error after some are found.
-    with contextlib.suppress(tokenize.TokenError):
-        try:
-            tokens = tokenize.tokenize(io.BytesIO(docstr.encode("utf-8")).readline)
-            for token in tokens:
-                state.add_token(token)
-        except IndentationError:
-            return None
+
+    # Keep tokenizing after an error. If `TokenError` is enountered, tokenize() will
+    # stop. We check the remaining bytes in bytes_io and resume tokenizing on the next
+    # loop iteration.
+    encoded_docstr = docstr.encode("utf-8")
+    bytes_io = io.BytesIO(encoded_docstr)
+    while bytes_io.tell() < len(encoded_docstr):
+        # Return all found signatures, even if there is a parse error after some are found.
+        with contextlib.suppress(tokenize.TokenError):
+            try:
+                tokens = tokenize.tokenize(bytes_io.readline)
+                for token in tokens:
+                    state.add_token(token)
+            except IndentationError:
+                return None
+
     sigs = state.get_signatures()
 
     def is_unique_args(sig: FunctionSig) -> bool:

mypy/test/teststubgen.py

@@ -429,6 +429,59 @@ def test_infer_sig_from_docstring_square_brackets(self) -> None:
             == []
         )
 
+    def test_infer_sig_from_docstring_unterminated_string_literal(self) -> None:
+        docstring = """
+        func(*args, **kwargs)
+        Overloaded function.
+
+        1. func(x: int) -> None
+
+        This is a valid docstring with an "unterminated string literal.
+
+        2. func(x: int, y: int) -> str
+
+        This is an overloaded method.
+        """
+        sigs = infer_sig_from_docstring(docstring, name="func")
+        assert_equal(
+            sigs[0], FunctionSig(name="func", args=[ArgSig(name="x", type="int")], ret_type="None")
+        )
+        assert_equal(
+            sigs[1],
+            FunctionSig(
+                name="func",
+                args=[ArgSig(name="x", type="int"), ArgSig(name="y", type="int")],
+                ret_type="str",
+            ),
+        )
+
+    def test_infer_sig_from_docstring_latex(self) -> None:
+        docstring = """
+        func(*args, **kwargs)
+        Overloaded function.
+
+        1. func(x: int) -> None
+
+        .. math::
+            \\mathbf{f}\\left(x\\right) = \\pi \\cdot x
+
+        2. func(x: int, y: int) -> str
+
+        This is an overloaded method.
+        """
+        sigs = infer_sig_from_docstring(docstring, name="func")
+        assert_equal(
+            sigs[0], FunctionSig(name="func", args=[ArgSig(name="x", type="int")], ret_type="None")
+        )
+        assert_equal(
+            sigs[1],
+            FunctionSig(
+                name="func",
+                args=[ArgSig(name="x", type="int"), ArgSig(name="y", type="int")],
+                ret_type="str",
+            ),
+        )
+
     def test_remove_misplaced_type_comments_1(self) -> None:
         good = """
         \u1234