Update API and license tools

utils/generate_api.py

@@ -1,9 +1,15 @@
 #!/usr/bin/env python
+# Licensed to Elasticsearch B.V under one or more agreements.
+# Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+# See the LICENSE file in the project root for more information
+
 
 import os
 import json
 import re
+import urllib3
 from itertools import chain
+from functools import lru_cache
 
 import black
 from click.testing import CliRunner
@@ -12,11 +18,14 @@ from pathlib import Path
 from jinja2 import Environment, FileSystemLoader, TemplateNotFound
 
 
+http = urllib3.PoolManager()
+
 # line to look for in the original source file
 SEPARATOR = "    # AUTO-GENERATED-API-DEFINITIONS #"
 # global substitutions for python keywords
 SUBSTITUTIONS = {"type": "doc_type", "from": "from_"}
 # api path(s)
+BRANCH_NAME = "7.x"
 CODE_ROOT = Path(__file__).absolute().parent.parent
 BASE_PATH = (
     CODE_ROOT.parent
@@ -53,6 +62,11 @@ def blacken(filename):
     assert result.exit_code == 0, result.output
 
 
+@lru_cache()
+def is_valid_url(url):
+    return http.request("HEAD", url).status == 200
+
+
 class Module:
     def __init__(self, namespace):
         self.namespace = namespace
@@ -135,6 +149,14 @@ class API:
             if self.doc_url.startswith("http://"):
                 self.doc_url = self.doc_url.replace("http://", "https://")
 
+            # Try setting doc refs like 'current' and 'master' to our branches ref.
+            if BRANCH_NAME is not None:
+                revised_url = re.sub("/elasticsearch/reference/[^/]+/", f"/elasticsearch/reference/{BRANCH_NAME}/", self.doc_url)
+                if is_valid_url(revised_url):
+                    self.doc_url = revised_url
+                else:
+                    print(f"URL {revised_url!r}, falling back on {self.doc_url!r}")
+
     @property
     def all_parts(self):
         parts = {}

utils/license_headers.py

@@ -0,0 +1,95 @@
+# Licensed to Elasticsearch B.V under one or more agreements.
+# Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+# See the LICENSE file in the project root for more information
+
+"""Script which verifies that all source files have a license header.
+Has two modes: 'fix' and 'check'. 'fix' fixes problems, 'check' will
+error out if 'fix' would have changed the file.
+"""
+
+import os
+import sys
+from typing import List, Iterator
+from itertools import chain
+
+
+lines_to_keep = [
+    "# -*- coding: utf-8 -*-\n",
+    "#!/usr/bin/env python\n"
+]
+license_header_lines = [
+    "# Licensed to Elasticsearch B.V under one or more agreements.\n",
+    "# Elasticsearch B.V licenses this file to you under the Apache 2.0 License.\n",
+    "# See the LICENSE file in the project root for more information\n",
+    "\n"
+]
+
+
+def find_files_to_fix(sources: List[str]) -> Iterator[str]:
+    """Iterates over all files and dirs in 'sources' and returns
+    only the filepaths that need fixing.
+    """
+    for source in sources:
+        if os.path.isfile(source) and does_file_need_fix(source):
+            yield source
+        elif os.path.isdir(source):
+            for root, _, filenames in os.walk(source):
+                for filename in filenames:
+                    filepath = os.path.join(root, filename)
+                    if does_file_need_fix(filepath):
+                        yield filepath
+
+
+def does_file_need_fix(filepath: str) -> bool:
+    if not filepath.endswith(".py"):
+        return False
+    with open(filepath, mode="r") as f:
+        first_license_line = None
+        for line in f:
+            if line == license_header_lines[0]:
+                first_license_line = line
+                break
+            elif line not in lines_to_keep:
+                return True
+        for header_line, line in zip(license_header_lines, chain((first_license_line,), f)):
+            if line != header_line:
+                return True
+    return False
+
+
+def add_header_to_file(filepath: str) -> None:
+    with open(filepath, mode="r") as f:
+        lines = list(f)
+    i = 0
+    for i, line in enumerate(lines):
+        if line not in lines_to_keep:
+            break
+    lines = lines[:i] + license_header_lines + lines[i:]
+    with open(filepath, mode="w") as f:
+        f.truncate()
+        f.write("".join(lines))
+    print(f"Fixed {os.path.relpath(filepath, os.getcwd())}")
+
+
+def main():
+    mode = sys.argv[1]
+    assert mode in ("fix", "check")
+    sources = [os.path.abspath(x) for x in sys.argv[2:]]
+    files_to_fix = find_files_to_fix(sources)
+
+    if mode == "fix":
+        for filepath in files_to_fix:
+            add_header_to_file(filepath)
+    else:
+        no_license_headers = list(files_to_fix)
+        if no_license_headers:
+            print("No license header found in:")
+            cwd = os.getcwd()
+            [print(f" - {os.path.relpath(filepath, cwd)}") for filepath in no_license_headers]
+            sys.exit(1)
+        else:
+            print("All files had license header")
+
+
+if __name__ == "__main__":
+    main()