import json

import os

import re

import sys

from pathlib import Path

from typing import TYPE_CHECKING

from invoke import task

from jsondiff import diff

if TYPE_CHECKING:

from invoke import Context, Result

TOP_DIR = Path(__file__).parent.resolve()

def update_file(

filename: Path | str, sub_line: tuple[str, str], strip: str | None = None

):

"""Utility function for tasks to read, update, and write files"""

with open(filename) as handle:

lines = [

re.sub(sub_line[0], sub_line[1], line.rstrip(strip)) for line in handle

]

with open(filename, "w") as handle:

handle.write("\n".join(lines))

handle.write("\n")

@task

def generate_openapi(_):

"""Update OpenAPI schema in file 'local_openapi.json'"""

from optimade.server.main import app

if not TOP_DIR.joinpath("openapi").exists():

os.mkdir(TOP_DIR.joinpath("openapi"))

with open(TOP_DIR.joinpath("openapi/local_openapi.json"), "w") as f:

json.dump(app.openapi(), f, indent=2)

print("", file=f) # Empty EOL

@task

def generate_index_openapi(_):

"""Update OpenAPI schema in file 'local_index_openapi.json'"""

from optimade.server.main_index import app as app_index

if not TOP_DIR.joinpath("openapi").exists():

os.mkdir(TOP_DIR.joinpath("openapi"))

with open(TOP_DIR.joinpath("openapi/local_index_openapi.json"), "w") as f:

json.dump(app_index.openapi(), f, indent=2)

print("", file=f) # Empty EOL

@task(pre=[generate_openapi, generate_index_openapi])

def check_openapi_diff(_):

"""Checks the Generated OpenAPI spec against what is stored in the repo"""

with open(TOP_DIR.joinpath("openapi/openapi.json")) as handle:

openapi = json.load(handle)

with open(TOP_DIR.joinpath("openapi/local_openapi.json")) as handle:

local_openapi = json.load(handle)

with open(TOP_DIR.joinpath("openapi/index_openapi.json")) as handle:

index_openapi = json.load(handle)

with open(TOP_DIR.joinpath("openapi/local_index_openapi.json")) as handle:

local_index_openapi = json.load(handle)

openapi_diff = diff(openapi, local_openapi)

if openapi_diff != {}:

print(

"Error: Generated OpenAPI spec for test server did not match committed version.\n"

"Run 'invoke update-openapijson' and re-commit.\n"

f"Diff:\n{openapi_diff}"

)

sys.exit(1)

openapi_index_diff = diff(index_openapi, local_index_openapi)

if openapi_index_diff != {}:

print(

"Error: Generated OpenAPI spec for Index meta-database did not match committed version.\n"

"Run 'invoke update-openapijson' and re-commit.\n"

f"Diff:\n{openapi_index_diff}"

)

sys.exit(1)

@task(pre=[generate_openapi, generate_index_openapi])

def update_openapijson(c):

"""Updates the stored OpenAPI spec to what the server returns"""

c.run(

f"cp {TOP_DIR.joinpath('openapi/local_openapi.json')} {TOP_DIR.joinpath('openapi/openapi.json')}"

)

c.run(

f"cp {TOP_DIR.joinpath('openapi/local_index_openapi.json')} {TOP_DIR.joinpath('openapi/index_openapi.json')}"

)

print("Updated OpenAPI JSON specifications")

@task(help={"ver": "OPTIMADE Python tools version to set"}, post=[update_openapijson])

def setver(_, ver=""):

"""Sets the OPTIMADE Python tools version"""

match = re.fullmatch(r"v?([0-9]+\.[0-9]+\.[0-9]+)", ver)

if not match or (match and len(match.groups()) != 1):

print(

"Error: Please specify version as 'Major.Minor.Patch' or 'vMajor.Minor.Patch'"

)

sys.exit(1)

ver = match.group(1)

update_file(

TOP_DIR.joinpath("optimade/__init__.py"),

(r'__version__ = ".*"', f'__version__ = "{ver}"'),

)

update_file(

TOP_DIR.joinpath("docs/static/default_config.json"),

(r'"version": ".*",', f'"version": "{ver}",'),

)

print(f"Bumped version to {ver}")

@task(help={"ver": "OPTIMADE API version to set"}, post=[update_openapijson])

def set_optimade_ver(_, ver=""):

"""Sets the OPTIMADE API Version"""

if not ver:

print("Error: Please specify --ver='Major.Minor.Patch(-rc|a|b.NUMBER)'")

sys.exit(1)

match = re.fullmatch(r"v?([0-9]+\.[0-9]+\.[0-9]+(-(rc|a|b)+\.[0-9]+)?)", ver)

if not match or (match and len(match.groups()) != 3):

print(

"Error: ver MUST be expressed as 'Major.Minor.Patch(-rc|a|b.NUMBER)'. It may be prefixed by 'v'."

)

sys.exit(1)

ver = match.group(1)

update_file(

TOP_DIR.joinpath("optimade/__init__.py"),

("__api_version__ = .+", f'__api_version__ = "{ver}"'),

)

for regex, version in (

("[0-9]+", ver.split(".")[0]),

("[0-9]+.[0-9]+", ".".join(ver.split(".")[:2])),

("[0-9]+.[0-9]+.[0-9]+", ver),

):

update_file(

TOP_DIR.joinpath("README.md"),

(f"example/v{regex}", f"example/v{version}"),

strip="\n",

)

update_file(

TOP_DIR.joinpath("optimade-version.json"),

(r'"message": "v.+",', f'"message": "v{ver}",'),

)

update_file(

TOP_DIR.joinpath("INSTALL.md"),

(r"/v[0-9]+(\.[0-9]+){2}(-(rc|a|b)+\.[0-9]+)?", f"/v{ver.split('.')[0]}"),

)

print(f"Bumped OPTIMADE API version to {ver}")

@task

def get_markdown_spec(ctx):

"""Convert the develop OPTIMADE specification from `rst` to `md`."""

print("Attempting to run pandoc...")

ctx.run(

"pandoc "

"-f rst -t markdown_strict "

"--wrap=preserve --columns=50000 "

"https://raw.githubusercontent.com/Materials-Consortia/OPTIMADE/develop/optimade.rst "

"> optimade.md"

)

@task(

help={

"pre-clean": "Remove the 'api_reference' sub directory prior to (re)creation."

}

)

def create_api_reference_docs(context, pre_clean=False, pre_commit=False):

"""Create the API Reference in the documentation"""

import shutil

def write_file(full_path: Path, content: str):

"""Write file with `content` to `full_path`"""

if full_path.exists():

with open(full_path) as handle:

cached_content = handle.read()

if content == cached_content:

del cached_content

return

del cached_content

with open(full_path, "w") as handle:

handle.write(content)

optimade_dir = TOP_DIR.joinpath("optimade")

docs_dir = TOP_DIR.joinpath("docs/api_reference")

unwanted_subdirs = ("__pycache__", "data", "grammar", "static")

pages_template = 'title: "{name}"\n'

md_template = "# {name}\n\n::: {py_path}\n"

models_template = (

md_template + f"{' ' * 4}options:\n{' ' * 6}show_if_no_docstring: true\n"

)

if docs_dir.exists() and pre_clean:

shutil.rmtree(docs_dir, ignore_errors=True)

if docs_dir.exists():

raise RuntimeError(f"{docs_dir} should have been removed!")

docs_dir.mkdir(exist_ok=True)

for dirpath, dirnames, filenames in os.walk(optimade_dir):

for unwanted_dir in unwanted_subdirs:

if unwanted_dir in dirnames:

# Avoid walking into or through unwanted directories

dirnames.remove(unwanted_dir)

relpath = Path(dirpath).relative_to(optimade_dir)

# Create `.pages`

if str(relpath) == ".":

write_file(

full_path=docs_dir.joinpath(".pages"),

content=pages_template.format(name="API Reference"),

)

docs_sub_dir = docs_dir.joinpath(relpath)

docs_sub_dir.mkdir(exist_ok=True)

if str(relpath) != ".":

write_file(

full_path=docs_sub_dir.joinpath(".pages"),

content=pages_template.format(name=str(relpath).split("/")[-1]),

)

# Create markdown files

for filename in filenames:

if re.match(r".*\.py$", filename) is None or filename == "__init__.py":

# Not a Python file: We don't care about it!

# Or filename is `__init__.py`: We don't want it!

continue

basename = filename[: -len(".py")]

py_path = f"optimade/{relpath}/{basename}".replace("/", ".")

if str(relpath) == ".":

py_path = py_path.replace("...", ".")

print(filename, basename, py_path)

md_filename = filename.replace(".py", ".md")

# For models we want to include EVERYTHING, even if it doesn't have a doc-string

# Also for `server.config`

template = (

models_template

if str(relpath) == "models"

or (str(relpath) == "server" and basename == "config")

else md_template

)

write_file(

full_path=docs_sub_dir.joinpath(md_filename),

content=template.format(name=basename, py_path=py_path),

)

if pre_commit:

# Check if there have been any changes.

# List changes if yes.

if TYPE_CHECKING:

context: "Context" = context # type: ignore[no-redef]

# NOTE: grep returns an exit code of 1 if it doesn't find anything

# (which will be good in this case).

# Concerning the weird last grep command see:

# http://manpages.ubuntu.com/manpages/precise/en/man1/git-status.1.html

result: "Result" = context.run(

"git status --porcelain docs/api_reference | grep -E '^[? MARC][?MD]' || exit 0",

hide=True,

)

if result.stdout:

sys.exit(

f"The following files have been changed/added, please stage them:\n\n{result.stdout}"

)

@task(help={"fname": "The JSON file containing the OpenAPI schema to validate"})

def swagger_validator(_, fname):

"""This task can be used in the CI to test the generated OpenAPI schemas

with the online swagger validator.

Returns:

Non-zero exit code if validation fails, otherwise returns `0`.

"""

import requests

def print_error(string):

for line in string.split("\n"):

print(f"\033[31m{line}\033[0m")

swagger_url = "https://validator.swagger.io/validator/debug"

with open(fname) as f:

schema = json.load(f)

response = requests.post(swagger_url, json=schema)

if response.status_code != 200:

print_error(f"Server returned status code {response.status_code}.")

sys.exit(1)

try:

json_response = response.json()

except json.JSONDecodeError:

print_error(f"Unable to parse validator response as JSON: {response}")

sys.exit(1)

if json_response:

if any(json_response[k] for k in json_response):

print_error(f"Schema file {fname} did not pass validation.\n")

print_error(json.dumps(response.json(), indent=2))

sys.exit(1)