dataunitylab/json-schema

1import argparse2import copy3import gzip4import json5import os6from pathlib import Path7import sys8 9import numpy as np10import pybktree11from sklearn.model_selection import GroupShuffleSplit12import tqdm13import unionfind14import Levenshtein15 16 17PERMISSIVE_LICENSES = set(json.load(open("permissive_licenses.json")))18 19 20def files_list(licenses):21    data_path = Path("valid_data")22    files = [23        f24        for f in data_path.rglob("*.json")25        if f.is_file() and licenses["/".join(f.parts[1:3])] in PERMISSIVE_LICENSES26    ]27    return files28 29 30def write_schemas(filename, schema_list, schema_data):31    sys.stderr.write(f"Writing {filename}…\n")32    with gzip.open(Path("data") / filename, "wt") as f:33        for schema in tqdm.tqdm(list(schema_list)):34            filename = str(os.path.join(*Path(schema).parts[1:]))35 36            # Skip schemas that have not been fetched this run37            try:38                data = schema_data[filename]39            except KeyError:40                continue41 42            schema = open(schema).read()43 44            # Get stars or null if missing45            try:46                repoStars = int(data["repoStars"])47            except (KeyError, ValueError):48                repoStars = None49 50            obj = {51                "repository": data["repository"],52                "commit": data["commit"],53                "commitDate": data["commitDate"],54                "path": data["path"],55                "repoStars": repoStars,56                "repoLastFetched": data["repoLastFetched"],57                "content": schema,58                "license": data["license"],59                "language": data["language"],60            }61            json.dump(obj, f)62            f.write("\n")63 64 65def get_repo_data(file, key):66    data = {}67    with open(file, "r") as f:68        for line in f:69            obj = json.loads(line)70            data[obj["repository"]] = obj[key]71 72    return data73 74 75def main(similarity, split, seed, commits_file, licenses_file, languages_file):76    licenses = get_repo_data(licenses_file, "license")77    languages = get_repo_data(languages_file, "language")78    files = files_list(licenses)79 80    # Prepare a BK Tree if we're doing similarity grouping81    if similarity:82        tree = pybktree.BKTree(83            lambda a, b: Levenshtein.distance(a, b) / max(len(a), len(b))84        )85 86    # Initialize a union-find data structure87    uf = unionfind.UnionFind()88 89    # Track the first schema added to each org so we can group them90    org_map = {}91 92    sys.stderr.write("Grouping by repository…\n")93    for schema_file in tqdm.tqdm(files):94        path_str = str(schema_file)95 96        # Get the organization name from the path97        org = schema_file.parts[1:3]98 99        uf.add(str(schema_file))100        if org not in org_map:101            # Track the first schema for this organization102            org_map[org] = str(schema_file)103        else:104            # Merge with the previous group if this105            # organization has been seen before106            uf.union(org_map[org], str(schema_file))107 108        # Add to the BK Tree109        if similarity:110            tree.add((str(schema_file), open(schema_file).read().strip()))111 112    del org_map113 114    # Optionally group together similar files115    if similarity:116        sys.stderr.write("Grouping similar files…\n")117        for schema_file in tqdm.tqdm(files):118            path_str = str(schema_file)119            data = open(schema_file).read().strip()120 121            # Find similar schemas for this schema and group them together122            for other_path, _ in tree.find(data, similarity):123                uf.union(path_str, other_path)124 125    # Produce a list of schemas and their associated groups126    all_schemas = list()127    schema_groups = list()128    for group, schemas in enumerate(uf.components()):129        all_schemas.extend(schemas)130        schema_groups.extend([group] * len(schemas))131 132    # Split the schemas into training and test133    all_schemas = np.array(all_schemas)134    schema_groups = np.array(schema_groups)135    gss = GroupShuffleSplit(n_splits=1, train_size=split, random_state=seed)136    (train_indexes, test_indexes) = next(gss.split(all_schemas, groups=schema_groups))137 138    test_schemas = all_schemas[test_indexes]139    test_groups = schema_groups[test_indexes]140    gss = GroupShuffleSplit(n_splits=1, train_size=0.5, random_state=seed)141    (test_indexes, val_indexes) = next(gss.split(test_schemas, groups=test_groups))142 143    schema_data = {}144    with open(commits_file) as f:145        for line in f:146            obj = json.loads(line)147            for commit in obj["commits"]:148                obj = copy.deepcopy(obj)149                filename = os.path.join(obj["repository"], commit["sha"], obj["path"])150                obj["commit"] = commit["sha"]151                obj["commitDate"] = commit["date"]152                obj["license"] = licenses[obj["repository"]]153                obj["language"] = languages.get(obj["repository"])154                schema_data[filename] = obj155 156    # Write the train and test sets157    write_schemas("train.jsonl.gz", all_schemas[train_indexes], schema_data)158    write_schemas("test.jsonl.gz", test_schemas[test_indexes], schema_data)159    write_schemas("validation.jsonl.gz", test_schemas[val_indexes], schema_data)160 161 162if __name__ == "__main__":163    parser = argparse.ArgumentParser()164    parser.add_argument("--similarity", default=None, type=float)165    parser.add_argument("--seed", default=38, type=int)166    parser.add_argument("--split", default=0.8, type=float)167    parser.add_argument("--commits_file", default="commits.json")168    parser.add_argument("--licenses_file", default="licenses.json")169    parser.add_argument("--languages_file", default="languages.json")170    args = parser.parse_args()171    main(172        args.similarity,173        args.split,174        args.seed,175        args.commits_file,176        args.licenses_file,177        args.languages_file,178    )179