Match testcheck logic

Author: ilevkivskyi

mypy/test/testcmdline.py

@@ -58,7 +58,7 @@ def test_python_cmdline(testcase: DataDrivenTestCase, step: int) -> None:
     with open(program_path, "w", encoding="utf8") as file:
         for s in testcase.input:
             file.write(f"{s}\n")
-    args = parse_args(normalize_devnull(testcase.input[step - 1]), step)
+    args = parse_args(normalize_devnull("\n".join(testcase.input)), step)
     custom_cwd = parse_cwd(testcase.input[1]) if len(testcase.input) > 1 else None
     args.append("--show-traceback")
     if "--error-summary" not in args:
@@ -122,19 +122,22 @@ def test_python_cmdline(testcase: DataDrivenTestCase, step: int) -> None:
         )
 
 
-def parse_args(line: str, step: int) -> list[str]:
-    """Parse the first line of the program for the command line.
+def parse_args(text: str, step: int) -> list[str]:
+    """Parse the program for the command line.
 
     This should have the form
 
-      # cmd: mypy <options>
+      # cmd[N]: mypy <options>
 
     For example:
 
-      # cmd: mypy pkg/
+      # cmd: mypy pkg/ or # cmd2: mypy pkg/
     """
-    step_str = "" if step == 1 else str(step)
-    m = re.match(f"# cmd{step_str}: mypy (.*)$", line)
+    m = re.search("# cmd: mypy (.*)$", text, flags=re.MULTILINE)
+    if step > 1:
+        alt_m = re.search(f"# cmd{step}: mypy (.*)$", text, flags=re.MULTILINE)
+        if alt_m is not None:
+            m = alt_m
     if not m:
         return []  # No args; mypy will spit out an error.
     return shlex.split(m.group(1))