#!/usr/bin/env python3

"""Script for applying a cache diff.

With some infrastructure, this can allow for distributing small cache diffs to users in

many cases instead of full cache artifacts.

Use diff-cache.py to generate a cache diff.

"""

from __future__ import annotations

import argparse

import os

import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from librt import base64

from librt.internal import ReadBuffer

from mypy.cache import CacheMeta

from mypy.defaults import SQLITE_NUM_SHARDS

from mypy.metastore import FilesystemMetadataStore, MetadataStore, SqliteMetadataStore

from mypy.util import json_dumps, json_loads

def make_cache(input_dir: str, sqlite: bool, num_shards: int = SQLITE_NUM_SHARDS) -> MetadataStore:

if sqlite:

return SqliteMetadataStore(input_dir, num_shards=num_shards)

else:

return FilesystemMetadataStore(input_dir)

def apply_diff(

cache_dir: str, diff_file: str, sqlite: bool = False, num_shards: int = SQLITE_NUM_SHARDS

) -> None:

cache = make_cache(cache_dir, sqlite, num_shards=num_shards)

with open(diff_file, "rb") as f:

diff = json_loads(f.read())

old_deps = json_loads(cache.read("@deps.meta.json"))

for file, data in diff.items():

if data is None:

cache.remove(file)

else:

if file.endswith(".ff"):

data_bytes = base64.b64decode(data)

else:

data_bytes = data.encode() if isinstance(data, str) else data

cache.write(file, data_bytes)

if file.endswith(".meta.ff") and "@deps" not in file:

buf = ReadBuffer(data_bytes[2:])

meta = CacheMeta.read(buf, data_file="")

assert meta is not None

old_deps["snapshot"][meta.id] = meta.hash

elif file.endswith(".meta.json") and "@deps" not in file:

meta_dict = json_loads(data_bytes)

old_deps["snapshot"][meta_dict["id"]] = meta_dict["hash"]

cache.write("@deps.meta.json", json_dumps(old_deps))

cache.commit()

def main() -> None:

parser = argparse.ArgumentParser()

parser.add_argument("--sqlite", action="store_true", default=False, help="Use a sqlite cache")

parser.add_argument(

"--num-shards", type=int, default=SQLITE_NUM_SHARDS, help=argparse.SUPPRESS

)

parser.add_argument("cache_dir", help="Directory for the cache")

parser.add_argument("diff", help="Cache diff file")

args = parser.parse_args()

apply_diff(args.cache_dir, args.diff, args.sqlite, num_shards=args.num_shards)

if __name__ == "__main__":

main()