Skip to content

Commit

Permalink
Expand string annotation for TypeAlias declarations
Browse files Browse the repository at this point in the history
Reviewed By: kinto0

Differential Revision: D48217714

fbshipit-source-id: ee48c2d625400a48cbc88a12d39d5ec0df45937a
  • Loading branch information
grievejia authored and facebook-github-bot committed Aug 14, 2023
1 parent e21527d commit e6613db
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 7 deletions.
34 changes: 27 additions & 7 deletions source/analysis/preprocessing.ml
Original file line number Diff line number Diff line change
Expand Up @@ -132,14 +132,34 @@ let transform_annotations ~transform_annotation_expression source =
let transform_children state _ = state, true

let transform_assign ~assign:({ Assign.annotation; value = assign_value; _ } as assign) =
match Node.value assign_value with
| Expression.Call { callee; _ } when is_type_variable_definition callee ->
match annotation with
| Some
{
assign with
Assign.annotation = annotation >>| transform_annotation_expression;
Assign.value = transform_annotation_expression assign_value;
}
| _ -> { assign with Assign.annotation = annotation >>| transform_annotation_expression }
Node.value =
Expression.Name
(Name.Attribute
{
base =
{
Node.value =
Expression.Name (Name.Identifier ("typing" | "typing_extensions"));
_;
};
attribute = "TypeAlias";
_;
});
_;
} ->
{ assign with Assign.value = transform_annotation_expression assign_value }
| _ -> (
match Node.value assign_value with
| Expression.Call { callee; _ } when is_type_variable_definition callee ->
{
assign with
Assign.annotation = annotation >>| transform_annotation_expression;
Assign.value = transform_annotation_expression assign_value;
}
| _ -> { assign with Assign.annotation = annotation >>| transform_annotation_expression })


let statement _ ({ Node.value; _ } as statement) =
Expand Down
7 changes: 7 additions & 0 deletions source/analysis/test/preprocessingTest.ml
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,13 @@ let test_expand_string_annotations _ =
{|
def foo() -> TakesParamSpec[[a, b]]: ...
|};
assert_expand
{|
X: typing.TypeAlias = "Dict[int, X] | int"
|}
{|
X: typing.TypeAlias = Dict[int, X] | int
|};

(* Ensure init subclass arguments are not counted as annotations to be expanded *)
assert_expand
Expand Down

0 comments on commit e6613db

Please sign in to comment.