From 0de26fc769d2af3b741d3d218348633a232f11b9 Mon Sep 17 00:00:00 2001 From: Jennifer Taylor Date: Tue, 15 Oct 2019 13:25:22 -0700 Subject: [PATCH] Address TODO in visitor type checking using mapping of CSTNode to valid return types. --- libcst/matchers/_visitors.py | 22 ++++++---------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/libcst/matchers/_visitors.py b/libcst/matchers/_visitors.py index 2203a159..07ee4928 100644 --- a/libcst/matchers/_visitors.py +++ b/libcst/matchers/_visitors.py @@ -37,6 +37,7 @@ from libcst.matchers._matcher_base import ( OneOf, matches, ) +from libcst.matchers._return_types import TYPED_FUNCTION_RETURN_MAPPING _CSTNodeT = TypeVar("_CSTNodeT", bound=cst.CSTNode) @@ -58,11 +59,11 @@ def _get_possible_match_classes(matcher: BaseMatcherNode) -> List[Type[cst.CSTNo return [getattr(cst, matcher.__class__.__name__)] -def _get_possible_annotated_classes(annotation: Type[object]) -> List[Type[object]]: +def _get_possible_annotated_classes(annotation: object) -> List[Type[object]]: if getattr(annotation, "__origin__", None) is Union: return getattr(annotation, "__args__", []) else: - return [annotation] + return [cast(Type[object], annotation)] def _get_valid_leave_annotations_for_classes( @@ -75,20 +76,9 @@ def _get_valid_leave_annotations_for_classes( # all possible valid return annotations. Its not really possible for us (or # pyre) to fully enforce return types given the presence of OneOf/AllOf matchers, so # we do the best we can by taking a union of all valid return annotations. - - # TODO: We could possibly teach LibCST codegen to generate a mapping of class - # to valid leave annotations when it generates leave_ methods, but for - # now this is functionally identical. That would get rid of the need to provide - # a namespace here, as well as a gross getattr on the CSTTransformer class. - meth = getattr(cst.CSTTransformer, f"leave_{cls.__name__}") - namespace: Dict[str, object] = { - **{x: getattr(cst, x) for x in dir(cst)}, - "cst": cst, - "libcst": cst, - } - type_hints = get_type_hints(meth, namespace, namespace) - if "return" in type_hints: - retval.update(_get_possible_annotated_classes(type_hints["return"])) + retval.update( + _get_possible_annotated_classes(TYPED_FUNCTION_RETURN_MAPPING[cls]) + ) return retval