DLSupplyChainLandscape/extract_dependencies.py at main · gaokai320/DLSupplyChainLandscape · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import json
from tqdm import tqdm
import logging
from pymongo import MongoClient
from collections import defaultdict
from packaging.requirements import Requirement, InvalidRequirement
from packaging.markers import UndefinedEnvironmentName

db = MongoClient(host="127.0.0.1", port=27017)['pypi']
distribution_metadata = db['distribution_metadata']
# print(distribution_metadata.count_documents({}))


def parse_requirement(requirement_str: str):
    try:
        req = Requirement(requirement_str)
    except InvalidRequirement:
        logging.error("InvalidRequirement: {}".format(requirement_str))
        return
    except:
        logging.error("ParseError: {}".format(requirement_str))
        return
    name, extras, specifier, marker = req.name, req.extras, req.specifier, req.marker
    if len(extras) > 0:
        name = name + '[' + ','.join(extras) + ']'
    specifier = [str(s) for s in specifier]
    extra = False
    if marker is not None:
        try:
            marker.evaluate()
        except UndefinedEnvironmentName:
            extra = True
        except:
            logging.error("MarkerError: {}".format(marker))
            extra = False
    return name, tuple(specifier), extra


def transform():
    dependencies = db['dependencies']
    dependencies.drop()
    dependencies = db['dependencies']
    deps = defaultdict(list)
    for doc in tqdm(distribution_metadata.find({}, {"name": 1, "version": 1, "_id": 0, "requires_dist": 1})):
        name = doc.get("name", "")
        version = doc.get('version', '')
        requires_dist = doc.get('requires_dist', [])
        key_name = name + ' ' + version
        for req in requires_dist:
            res = parse_requirement(req)
            if res is not None:
                dependency, dependency_version, extra = res
                deps[key_name].append((dependency, dependency_version, extra))
    print('Dumping to json file')
    with open('/fast/pypi/dependency_dump.json', 'w') as outf:
        json.dump(deps, outf)
    transformed_docs = []
    print('Inserting')
    for k, v in tqdm(deps.items()):
        name, version = k.split(' ', 1)
        for dependency, dependency_version, extra in set(v):
            transformed_docs.append({"name": name, "version": version,
                                     "dependency": dependency, "dependency_version": dependency_version,
                                     "extra": extra})
    print(len(transformed_docs))
    dependencies.insert_many(transformed_docs)
    print('Finished')


if __name__ == "__main__":
    # req = 'psycopg2[abc, def] (>=2.8) ; (platform_python_implementation != "PyPy") and extra=="postgres"'
    # print(parse_req(req))
    logging.basicConfig(format="%(asctime)s [%(levelname)s] %(message)s", level=logging.INFO,
                        filename='log/extract_dependency.log', filemode='w')
    transform()
    # test = [
    #     'psycopg2[abc, def] (>=2.8) ; (platform_python_implementation != "PyPy") and extra=="postgres"',
    #     "numpy",
    #     "sphinx-rtd-theme (==0.5.0) ; extra == 'all'",
    #     "sphinx-rtd-theme (=0.5.0) ; extra == 'all'",
    #     "django-configurations[database,email] ; extr == 'configuration'",
    #     "prompt-toolkit (<=2.0.10,>=2.0.4)",
    #     "importlib-resources ; python_version < '3.7'"
    # ]
    # for t in test:
    #     res = parse_requirement(t)
    #     if res is not None:
    #         name, specs, extra = res
    #         print(name, specs, extra)