From a785041250910bbf8711c224b94c268c67e1c683 Mon Sep 17 00:00:00 2001
From: Akuli <akuviljanen17@gmail.com>
Date: Wed, 30 Aug 2023 17:43:46 +0300
Subject: [PATCH] create_baseline_stubs.py: Improve pyright config file editing
 (#10629)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 pyrightconfig.stricter.json      | 14 ++++++------
 scripts/create_baseline_stubs.py | 38 +++++++++++++++++++++++---------
 2 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/pyrightconfig.stricter.json b/pyrightconfig.stricter.json
index fb5792305..e3a3839a0 100644
--- a/pyrightconfig.stricter.json
+++ b/pyrightconfig.stricter.json
@@ -23,10 +23,10 @@
         "stdlib/xml/dom/pulldom.pyi",
         "stdlib/xml/sax",
         "stubs/aws-xray-sdk",
-        "stubs/bleach",
-        "stubs/boto",
         "stubs/beautifulsoup4",
+        "stubs/bleach",
         "stubs/boltons",
+        "stubs/boto",
         "stubs/braintree",
         "stubs/caldav",
         "stubs/cffi",
@@ -36,9 +36,11 @@
         "stubs/docutils",
         "stubs/Flask-Migrate",
         "stubs/fpdf2",
+        "stubs/google-cloud-ndb",
         "stubs/html5lib",
         "stubs/httplib2",
         "stubs/humanfriendly",
+        "stubs/influxdb-client",
         "stubs/invoke",
         "stubs/jmespath",
         "stubs/jsonschema",
@@ -47,14 +49,12 @@
         "stubs/mysqlclient",
         "stubs/oauthlib",
         "stubs/openpyxl",
-        "stubs/Pillow",
-        "stubs/protobuf",
-        "stubs/google-cloud-ndb",
-        "stubs/influxdb-client",
         "stubs/passlib",
         "stubs/peewee",
         "stubs/pexpect",
         "stubs/pika",
+        "stubs/Pillow",
+        "stubs/protobuf",
         "stubs/psutil",
         "stubs/psycopg2",
         "stubs/pyasn1",
@@ -75,7 +75,7 @@
         "stubs/urllib3",
         "stubs/vobject",
         "stubs/WebOb",
-        "stubs/workalendar"
+        "stubs/workalendar",
     ],
     "typeCheckingMode": "strict",
     // TODO: Complete incomplete stubs
diff --git a/scripts/create_baseline_stubs.py b/scripts/create_baseline_stubs.py
index e3c259e63..47cbafa35 100755
--- a/scripts/create_baseline_stubs.py
+++ b/scripts/create_baseline_stubs.py
@@ -152,22 +152,38 @@ def add_pyright_exclusion(stub_dir: str) -> None:
     assert i < len(lines), f"Error parsing {PYRIGHT_CONFIG}"
     while not lines[i].strip().startswith("]"):
         i += 1
-    # Must use forward slash in the .json file
-    line_to_add = f'        "{stub_dir}",'.replace("\\", "/")
-    initial = i - 1
-    while lines[i].lower() > line_to_add.lower():
+    end = i
+
+    # We assume that all third-party excludes must be at the end of the list.
+    # This helps with skipping special entries, such as "stubs/**/@tests/test_cases".
+    while lines[i - 1].strip().startswith('"stubs/'):
         i -= 1
-    if lines[i + 1].strip().rstrip(",") == line_to_add.strip().rstrip(","):
+    start = i
+
+    before_third_party_excludes = lines[:start]
+    third_party_excludes = lines[start:end]
+    after_third_party_excludes = lines[end:]
+
+    last_line = third_party_excludes[-1].rstrip()
+    if not last_line.endswith(","):
+        last_line += ","
+        third_party_excludes[-1] = last_line + "\n"
+
+    # Must use forward slash in the .json file
+    line_to_add = f'        "{stub_dir}",\n'.replace("\\", "/")
+
+    if line_to_add in third_party_excludes:
         print(f"{PYRIGHT_CONFIG} already up-to-date")
         return
-    if i == initial:
-        # Special case: when adding to the end of the list, commas need tweaking
-        line_to_add = line_to_add.rstrip(",")
-        lines[i] = lines[i].rstrip() + ",\n"
-    lines.insert(i + 1, line_to_add + "\n")
+
+    third_party_excludes.append(line_to_add)
+    third_party_excludes.sort(key=str.lower)
+
     print(f"Updating {PYRIGHT_CONFIG}")
     with open(PYRIGHT_CONFIG, "w", encoding="UTF-8") as f:
-        f.writelines(lines)
+        f.writelines(before_third_party_excludes)
+        f.writelines(third_party_excludes)
+        f.writelines(after_third_party_excludes)
 
 
 def main() -> None: