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
|
"""Main entry point to mypy annotation inference utility."""
import json
from typing import List
from mypy_extensions import TypedDict
from pyannotate_tools.annotations.types import ARG_STAR, ARG_STARSTAR
from pyannotate_tools.annotations.infer import infer_annotation
from pyannotate_tools.annotations.parse import parse_json
# Schema of a function signature in the output
Signature = TypedDict('Signature', {'arg_types': List[str],
'return_type': str})
# Schema of a function in the output
FunctionData = TypedDict('FunctionData', {'path': str,
'line': int,
'func_name': str,
'signature': Signature,
'samples': int})
SIMPLE_TYPES = {'None', 'int', 'float', 'str', 'bytes', 'bool'}
def unify_type_comments(type_comments):
# type: (List[str]) -> Signature
arg_types, return_type = infer_annotation(type_comments)
arg_strs = []
for arg, kind in arg_types:
arg_str = str(arg)
if kind == ARG_STAR:
arg_str = '*%s' % arg_str
elif kind == ARG_STARSTAR:
arg_str = '**%s' % arg_str
arg_strs.append(arg_str)
return {
'arg_types': arg_strs,
'return_type': str(return_type),
}
def is_signature_simple(signature):
# type: (Signature) -> bool
return (all(x.lstrip('*') in SIMPLE_TYPES for x in signature['arg_types']) and
signature['return_type'] in SIMPLE_TYPES)
def generate_annotations_json_string(source_path, only_simple=False):
# type: (str, bool) -> List[FunctionData]
"""Produce annotation data JSON file from a JSON file with runtime-collected types.
Data formats:
* The source JSON is a list of pyannotate_tools.annotations.parse.RawEntry items.
* The output JSON is a list of FunctionData items.
"""
items = parse_json(source_path)
results = []
for item in items:
signature = unify_type_comments(item.type_comments)
if is_signature_simple(signature) or not only_simple:
data = {
'path': item.path,
'line': item.line,
'func_name': item.func_name,
'signature': signature,
'samples': item.samples
} # type: FunctionData
results.append(data)
return results
def generate_annotations_json(source_path, target_path, only_simple=False):
# type: (str, str, bool) -> None
"""Like generate_annotations_json_string() but writes JSON to a file."""
results = generate_annotations_json_string(source_path, only_simple=only_simple)
with open(target_path, 'w') as f:
json.dump(results, f, sort_keys=True, indent=4)
|
