Skip to content

Commit 0e07c47

Browse files
tochilinaktochilinak
committed
first version of mypy annotation serializer
1 parent 7a0dc87 commit 0e07c47

5 files changed

Lines changed: 192 additions & 61 deletions

File tree

‎utbot-python/dev_test_samples/annotations_tests.py‎

Lines changed: 0 additions & 13 deletions
This file was deleted.

‎utbot-python/samples/easy_samples/annotations_tests.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
from typing import *
2+
import collections
3+
from enum import Enum
4+
5+
6+
def square(collection: Iterable[int], x: Any):
7+
result = set()
8+
for elem in collection:
9+
result.add(elem ** 2)
10+
return result
11+
12+
13+
def not_annotated(x):
14+
return x
15+
16+
17+
def same_annotations(x: int, y: int, a: List[Any], b: List[Any], c: List[int]):
18+
return x + y
19+
20+
21+
def optional(x: Optional[int]):
22+
return x
23+
24+
25+
def literal(x: Literal["w", "r"]):
26+
return x
27+
28+
29+
class Color(Enum):
30+
RED = 1
31+
GREEN = 2
32+
BLUE = 3
33+
34+
35+
def enum_literal(x: Literal[Color.RED, Color.GREEN]):
36+
return x
37+
38+
39+
if __name__ == "__main__":
40+
# print(square(collections.defaultdict(int)))
41+
# enum_literal(Color.BLUE)
42+
pass

‎utbot-python/src/main/kotlin/org/utbot/python/typing/PythonAnnotations.kt‎

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,9 @@ private val moshi = Moshi.Builder()
99
.add(
1010
PolymorphicJsonAdapterFactory.of(PythonAnnotationNode::class.java, "type")
1111
.withSubtype(ConcreteAnnotation::class.java, AnnotationType.Concrete.name)
12-
.withSubtype(AbstractAnnotation::class.java, AnnotationType.Abstract.name)
12+
.withSubtype(Protocol::class.java, AnnotationType.Protocol.name)
1313
.withSubtype(PythonAny::class.java, AnnotationType.Any.name)
1414
.withSubtype(PythonLiteral::class.java, AnnotationType.Literal.name)
15-
.withSubtype(PythonOptional::class.java, AnnotationType.Optional.name)
1615
.withSubtype(PythonUnion::class.java, AnnotationType.Union.name)
1716
.withSubtype(PythonList::class.java, AnnotationType.List.name)
1817
.withSubtype(PythonDict::class.java, AnnotationType.Dict.name)
@@ -25,10 +24,9 @@ private val jsonAdapter = moshi.adapter(PythonAnnotationStorage::class.java)
2524

2625
enum class AnnotationType {
2726
Concrete,
28-
Abstract,
27+
Protocol,
2928
Any,
3029
Literal,
31-
Optional,
3230
Union,
3331
List,
3432
Dict,
@@ -69,11 +67,11 @@ open class ConcreteAnnotation(
6967
get() = PythonClassId("$module.$simpleName")
7068
}
7169

72-
open class AbstractAnnotation(
70+
open class Protocol(
7371
val methods: List<String>,
7472
module: String,
7573
simpleName: String,
76-
type: AnnotationType = AnnotationType.Abstract,
74+
type: AnnotationType = AnnotationType.Protocol,
7775
): PythonAnnotationNode(module, simpleName, type)
7876

7977
open class SpecialAnnotation(
@@ -93,7 +91,6 @@ class PythonDict: GenericPrimitive("builtins", "dict", AnnotationType.Dict)
9391
class PythonTuple: GenericPrimitive("builtins", "tuple", AnnotationType.Tuple)
9492
class PythonAny: SpecialAnnotation("typing", "Any", AnnotationType.Any)
9593
class PythonLiteral: SpecialAnnotation("typing", "Literal", AnnotationType.Literal)
96-
class PythonOptional: SpecialAnnotation("typing", "Optional", AnnotationType.Optional)
9794
class PythonUnion: SpecialAnnotation("typing", "Union", AnnotationType.Union)
9895

9996
fun main() {
@@ -117,19 +114,23 @@ fun main() {
117114
{
118115
"nodeStorage": {
119116
"1": {
120-
"type": "Optional"
117+
"type": "Union"
121118
},
122119
"2": {
123120
"type": "Concrete",
124121
"module": "builtins",
125122
"simpleName": "int"
123+
},
124+
"3": {
125+
"type": "Dict"
126126
}
127127
},
128128
"annotations": [
129129
{
130130
"nodeId": "1",
131131
"args": [
132-
{ "nodeId": "2" }
132+
{ "nodeId": "2" },
133+
{ "nodeId": "3" }
133134
]
134135
}
135136
]
@@ -158,4 +159,13 @@ fun main() {
158159
val obj3 = jsonAdapter.fromJson(test3)!!
159160
println(obj3)
160161
println(obj3.normalizedRepr(obj3.annotations[0]))
162+
163+
val test4 = """
164+
{"nodeStorage": {"140503148502752": {"type": "Concrete", "module": "builtins", "simpleName": "int"}, "140503148586752": {"type": "List"}, "140503162870576": {"type": "Any"}, "140503162953056": {"type": "Any"}}, "annotations": [{"nodeId": "140503148502752"}, {"nodeId": "140503148502752"}, {"nodeId": "140503148586752", "args": [{"nodeId": "140503162953056"}]}, {"nodeId": "140503148586752", "args": [{"nodeId": "140503162870576"}]}, {"nodeId": "140503148586752", "args": [{"nodeId": "140503148502752"}]}]}
165+
""".trimIndent()
166+
val obj4 = jsonAdapter.fromJson(test4)!!
167+
println(obj4)
168+
obj4.annotations.forEach {
169+
println(obj4.normalizedRepr(it))
170+
}
161171
}

‎utbot-python/src/main/resources/mypy_main.py‎

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
11
from mypy.main import *
22

3+
import sys
4+
5+
from io import StringIO
6+
from typing import List, Tuple, TextIO, Callable
7+
38
"""
49
Copy with some changes of function 'main' from here:
510
https://github.com/python/mypy/blob/v0.971/mypy/main.py
@@ -50,7 +55,7 @@ def new_main(script_path: Optional[str],
5055

5156
if options.install_types and not sources:
5257
install_types(formatter, options, non_interactive=options.non_interactive)
53-
return
58+
return None
5459

5560
res, messages, blockers = run_build(sources, options, fscache, t0, stdout, stderr)
5661

@@ -98,7 +103,7 @@ def new_main(script_path: Optional[str],
98103
Copy with some changes of mypy api functions from here:
99104
https://github.com/python/mypy/blob/v0.971/mypy/api.py
100105
"""
101-
def _run(main_wrapper: Callable[[TextIO, TextIO], None]) -> Tuple[str, str, int, Optional[build.BuildResult]]:
106+
def _run(main_wrapper: Callable[[TextIO, TextIO], Optional[build.BuildResult]]) -> Tuple[str, str, int, Optional[build.BuildResult]]:
102107

103108
stdout = StringIO()
104109
stderr = StringIO()
@@ -113,6 +118,21 @@ def _run(main_wrapper: Callable[[TextIO, TextIO], None]) -> Tuple[str, str, int,
113118
return stdout.getvalue(), stderr.getvalue(), exit_status, res
114119

115120

116-
def run(args: List[str]) -> Tuple[str, str, int]:
121+
def run(args: List[str]) -> Tuple[str, str, int, Optional[build.BuildResult]]:
122+
args.append("--no-incremental")
117123
return _run(lambda stdout, stderr: new_main(None, args=args,
118124
stdout=stdout, stderr=stderr, clean_exit=True))
125+
126+
127+
if __name__ == "__main__":
128+
import time
129+
start = time.time()
130+
stdout, stderr, exit_status, build_result = run(sys.argv[1:])
131+
print(f"Seconds passed: {time.time() - start}")
132+
print(stdout, stderr, exit_status, sep='\n')
133+
if build_result is None:
134+
print("BuildResult is None")
135+
else:
136+
print(build_result.files.keys())
137+
# print(build_result.files['general'].names)
138+

‎utbot-python/src/main/resources/new_normalize_annotation_from_project.py‎

Lines changed: 108 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,60 +1,132 @@
1-
import mypy.fastparse
1+
import os
2+
import sys
3+
import json
24
import typing
35

6+
import mypy_main
7+
import mypy.nodes
8+
import mypy.types
49

5-
class __Annotation:
6-
def __init__(self, annotation_type):
7-
self.type = annotation_type
810

9-
def encode(self):
10-
return {"type": self.type}
11+
annotation_node_set = set()
1112

1213

13-
class __GenericAnnotation(__Annotation):
14-
def __init__(self, annotation_type, args):
15-
super(__GenericAnnotation, self).__init__(annotation_type)
16-
self.args = args
14+
class AnnotationNode:
15+
def __init__(self, annotation_type, id_):
16+
self.type = annotation_type
17+
self.id_ = id_
1718

1819
def encode(self):
19-
return super(__GenericAnnotation, self).encode() + {"args": self.args.encode()}
20-
21-
22-
class __ConcreteAnnotation(__Annotation):
23-
def __init__(self):
24-
super(__ConcreteAnnotation, self).__init__("Concrete")
25-
26-
27-
class __Any(__Annotation):
28-
def __init__(self):
29-
super(__Any, self).__init__("Any")
30-
20+
return {"type": self.type}
3121

32-
class __Literal(__GenericAnnotation):
33-
def __init__(self):
34-
super(__Literal, self).__init__("Literal")
22+
def __eq__(self, other):
23+
return self.id_ == other.id_
3524

25+
def __hash__(self):
26+
return hash(self.id_)
3627

37-
class __Optional(__GenericAnnotation):
38-
def __init__(self):
39-
super(__Optional, self).__init__("Optional")
4028

29+
class CustomAnnotationNode(AnnotationNode):
30+
def __init__(self, symbol_node):
31+
if isinstance(symbol_node, mypy.nodes.TypeInfo):
32+
if symbol_node.is_protocol:
33+
annotation_type = "Protocol"
34+
else:
35+
annotation_type = "Concrete"
36+
super().__init__(annotation_type, id(symbol_node))
37+
self.module = symbol_node.module_name
38+
self.simple_name = symbol_node._fullname.split('.')[-1]
39+
else:
40+
assert False, "Some SymbolNode wasn't considered"
4141

42-
class __Union(__GenericAnnotation):
43-
def __init__(self):
44-
super(__Union, self).__init__("Union")
42+
def encode(self):
43+
return dict(super().encode(), **{"module": self.module, "simpleName": self.simple_name})
4544

4645

47-
// TODO: edit mypy_path
48-
def main(function_name, source_path):
46+
class Annotation:
47+
def __init__(self, node_id: int, args: typing.Optional[typing.List['Annotation']] = None):
48+
self.node_id = node_id
49+
self.args = args
4950

51+
def encode(self):
52+
result = {"nodeId": str(self.node_id)}
53+
if self.args is not None:
54+
result["args"] = list(map(lambda x: x.encode(), self.args))
55+
return result
56+
57+
58+
def get_annotation(mypy_type) -> Annotation:
59+
if isinstance(mypy_type, mypy.types.Instance):
60+
children = []
61+
for arg in mypy_type.args:
62+
children.append(get_annotation(arg))
63+
64+
if len(children) == 0:
65+
children = None
66+
67+
if mypy_type.type._fullname == "builtins.list":
68+
cur_node = AnnotationNode("List", id(mypy_type.type))
69+
elif mypy_type.type._fullname == "builtins.dict":
70+
cur_node = AnnotationNode("Dict", id(mypy_type.type))
71+
else:
72+
cur_node = CustomAnnotationNode(mypy_type.type) # mypy_type.type: mypy.nodes.TypeInfo
73+
74+
annotation_node_set.add(cur_node)
75+
return Annotation(cur_node.id_, children)
76+
77+
elif isinstance(mypy_type, mypy.types.AnyType):
78+
cur_node = AnnotationNode("Any", id(mypy_type))
79+
annotation_node_set.add(cur_node)
80+
return Annotation(cur_node.id_)
81+
82+
# elif isinstance(mypy_type, mypy.types.LiteralType): # TODO: consider enum literals
83+
# return Annotation("Literal", id(mypy_type))
84+
85+
elif isinstance(mypy_type, mypy.types.UnionType):
86+
children = []
87+
for arg in mypy_type.items:
88+
children.append(get_annotation(arg))
89+
cur_node = AnnotationNode("Union", id(mypy_type))
90+
annotation_node_set.add(cur_node)
91+
return Annotation(cur_node.id_, children)
92+
93+
assert False, "Some annotation type wasn't considered"
94+
95+
96+
def get_output_json(annotations: typing.List[Annotation]):
97+
result = {}
98+
result['nodeStorage'] = {}
99+
for node in annotation_node_set:
100+
result['nodeStorage'][node.id_] = node.encode()
101+
result['annotations'] = []
102+
for annotation in annotations:
103+
result['annotations'].append(annotation.encode())
104+
return json.dumps(result)
105+
106+
107+
def main(function_name, source_path, mypy_config_file):
108+
stdout, stderr, exit_status, build_result = mypy_main.run([
109+
source_path, "--config-file", mypy_config_file])
110+
module_name = os.path.basename(source_path)[:-3]
111+
function_info = build_result.files[module_name].names[function_name].node
112+
113+
if function_info.type is None:
114+
sys.stderr.write("No annotation")
115+
exit(1)
116+
117+
arg_annotations = function_info.type.arg_types
118+
annotation_list = []
119+
for x in arg_annotations:
120+
annotation_list.append(get_annotation(x))
121+
122+
print(get_output_json(annotation_list))
50123

51124

52125
def get_args():
53126
function_name = sys.argv[1]
54127
source_path = sys.argv[2]
55-
for extra_path in sys.argv[3:]:
56-
sys.path.append(extra_path)
57-
return annotation, cur_module, path
128+
mypy_config_file = sys.argv[3]
129+
return function_name, source_path, mypy_config_file
58130

59131

60132
if __name__ == '__main__':

0 commit comments

Comments
 (0)