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
|
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from typing import List
from libcst.codegen.gather import imports, nodebases, nodeuses
generated_code: List[str] = []
generated_code.append("# Copyright (c) Meta Platforms, Inc. and affiliates.")
generated_code.append("#")
generated_code.append(
"# This source code is licensed under the MIT license found in the"
)
generated_code.append("# LICENSE file in the root directory of this source tree.")
generated_code.append("")
generated_code.append("")
generated_code.append("# This file was generated by libcst.codegen.gen_type_mapping")
generated_code.append("from typing import Dict as TypingDict, Type, Union")
generated_code.append("")
generated_code.append("from libcst._maybe_sentinel import MaybeSentinel")
generated_code.append("from libcst._removal_sentinel import RemovalSentinel")
generated_code.append("from libcst._nodes.base import CSTNode")
# Import the types we use. These have to be type guarded since it would
# cause an import cycle otherwise.
generated_code.append("")
generated_code.append("")
for module, objects in imports.items():
generated_code.append(f"from {module} import (")
generated_code.append(f" {', '.join(sorted(objects))}")
generated_code.append(")")
# Generate the base visit_ methods
generated_code.append("")
generated_code.append("")
generated_code.append(
"TYPED_FUNCTION_RETURN_MAPPING: TypingDict[Type[CSTNode], object] = {"
)
for node in sorted(nodebases.keys(), key=lambda node: node.__name__):
name = node.__name__
if name.startswith("Base"):
continue
valid_return_types: List[str] = [nodebases[node].__name__]
node_uses = nodeuses[node]
base_uses = nodeuses[nodebases[node]]
if node_uses.maybe or base_uses.maybe:
valid_return_types.append("MaybeSentinel")
if (
node_uses.optional
or node_uses.sequence
or base_uses.optional
or base_uses.sequence
):
valid_return_types.append("RemovalSentinel")
generated_code.append(f' {name}: Union[{", ".join(valid_return_types)}],')
generated_code.append("}")
if __name__ == "__main__":
# Output the code
print("\n".join(generated_code))
|
