scripts: add diff output to export_requirements

Author: onerandomusername

scripts/_utils.py

@@ -0,0 +1,85 @@
+"""Utility functions and variables which are useful for all scripts."""
+import difflib
+import importlib.util
+import os
+import pathlib
+import typing
+
+
+MODMAIL_DIR = pathlib.Path(importlib.util.find_spec("modmail").origin).parent
+PROJECT_DIR = MODMAIL_DIR.parent
+try:
+    import pygments
+except ModuleNotFoundError:
+    pygments = None
+else:
+    from pygments.formatters import Terminal256Formatter
+    from pygments.lexers.diff import DiffLexer
+
+
+class CheckFileEdit:
+    """Check if a file is edited within the body of this class."""
+
+    def __init__(self, *files: os.PathLike):
+        self.files: typing.List[pathlib.Path] = []
+        for f in files:
+            self.files.append(pathlib.Path(f))
+        self.return_value: typing.Optional[int] = None
+        self.edited_files: typing.Dict[pathlib.Path] = dict()
+
+    def __enter__(self):
+        self.file_contents = {}
+        for file in self.files:
+            try:
+                with open(file, "r") as f:
+                    self.file_contents[file] = f.readlines()
+            except FileNotFoundError:
+                self.file_contents[file] = None
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):  # noqa: ANN001
+        for file in self.files:
+            with open(file, "r") as f:
+                original_contents = self.file_contents[file]
+                new_contents = f.readlines()
+                if original_contents != new_contents:
+                    # construct a diff
+                    diff = difflib.unified_diff(
+                        original_contents, new_contents, fromfile="before", tofile="after"
+                    )
+                    try:
+                        diff = "".join(diff)
+                    except TypeError:
+                        diff = None
+                    else:
+                        if pygments is not None:
+                            diff = pygments.highlight(diff, DiffLexer(), Terminal256Formatter())
+                    self.edited_files[file] = diff
+
+    def write(self, path: str, contents: typing.Union[str, bytes], *, force: bool = False, **kwargs) -> bool:
+        """
+        Write to the provided path with contents. Must be within the context manager.
+
+        Returns False if contents are not edited, True if they are.
+        If force is True, will modify the files even if the contents match.
+
+        Any extras kwargs are passed to open()
+        """
+        path = pathlib.Path(path)
+        if path not in self.files:
+            raise AssertionError(f"{path} must have been passed to __init__")
+
+        if not force:
+            try:
+                with open(path, "r") as f:
+                    if contents == f.read():
+                        return False
+            except FileNotFoundError:
+                pass
+        if isinstance(contents, str):
+            contents = contents.encode()
+
+        with open(path, "wb") as f:
+            f.write(contents)
+
+        return True

scripts/export_requirements.py

@@ -16,10 +16,12 @@
 
 import tomli
 
+from ._utils import PROJECT_DIR, CheckFileEdit
 
-GENERATED_FILE = pathlib.Path("requirements.txt")
-CONSTRAINTS_FILE = pathlib.Path("modmail/constraints.txt")
-DOC_REQUIREMENTS = pathlib.Path("docs/.requirements.txt")
+
+GENERATED_FILE = PROJECT_DIR / "requirements.txt"
+CONSTRAINTS_FILE = PROJECT_DIR / "modmail/constraints.txt"
+DOC_REQUIREMENTS = PROJECT_DIR / "docs/.requirements.txt"
 
 VERSION_RESTRICTER_REGEX = re.compile(r"(?P<sign>[<>=!]{1,2})(?P<version>\d+\.\d+?)(?P<patch>\.\d+?|\.\*)?")
 PLATFORM_MARKERS_REGEX = re.compile(r'sys_platform\s?==\s?"(?P<platform>\w+)"')
@@ -128,6 +130,7 @@ def _export_doc_requirements(toml: dict, file: pathlib.Path, *packages) -> int:
     file = pathlib.Path(file)
     if not file.exists():
         # file does not exist
+        print(f"{file.relative_to(PROJECT_DIR)!s} must exist to export doc requirements")
         return 2
 
     with open(file) as f:
@@ -149,14 +152,18 @@ def _export_doc_requirements(toml: dict, file: pathlib.Path, *packages) -> int:
     except AttributeError as e:
         print(e)
         return 3
-    if new_contents == contents:
-        # don't write anything, just return 0
-        return 0
 
-    with open(file, "w") as f:
-        f.write(new_contents)
+    with CheckFileEdit(file) as check_file:
+
+        check_file.write(file, new_contents)
 
-    return 1
+    for file, diff in check_file.edited_files.items():
+        print(
+            f"Exported new documentation requirements to {file.relative_to(PROJECT_DIR)!s}.",
+            file=sys.stderr,
+        )
+        print(diff or "No diff to show.")
+        print()
 
 
 def export(
@@ -269,19 +276,17 @@ def export(
     else:
         exit_code = 0
 
-    if req_path.exists():
-        with open(req_path, "r") as f:
-            if req_txt == f.read():
-                # nothing to edit
-                # if exit_code is ever removed from here, this should return zero
-                return exit_code
+    with CheckFileEdit(req_path) as check_file:
+        check_file.write(req_path, req_txt)
 
-    if _write_file(req_path, req_txt):
-        print(f"Updated {req_path} with new requirements.")
-        return 1
-    else:
-        print(f"No changes were made to {req_path}")
-        return 0
+    for file, diff in check_file.edited_files.items():
+        print(
+            f"Exported new requirements to {file.relative_to(PROJECT_DIR)}.",
+            file=sys.stderr,
+        )
+        print(diff or "No diff to show.")
+        print()
+    return bool(len(check_file.edited_files)) or exit_code
 
 
 def main(path: os.PathLike, include_markers: bool = True, **kwargs) -> int: