From d9436e276fbab946c53a947904af99b48c5612e2 Mon Sep 17 00:00:00 2001 From: Cheukting Date: Wed, 11 May 2022 12:13:59 +0100 Subject: [PATCH 1/8] flatten literal --- src/shed/_codemods.py | 30 ++++++++++++++++++++++++++++++ tests/recorded/flatten_literal.txt | 7 +++++++ 2 files changed, 37 insertions(+) create mode 100644 tests/recorded/flatten_literal.txt diff --git a/src/shed/_codemods.py b/src/shed/_codemods.py index 6e0219e..192771e 100644 --- a/src/shed/_codemods.py +++ b/src/shed/_codemods.py @@ -332,3 +332,33 @@ def _has_none(node): return updated_node.with_changes(left=updated_node.right, right=node_left) else: return updated_node + + @m.call_if_inside(m.Annotation(annotation=m.BinaryOperation())) + @m.leave( + m.BinaryOperation( + left=m.Subscript(value=m.Name(value="Literal")), + operator=m.BitOr(), + right=m.Name("None"), + ) + ) + def flatten_literal(self, _, updated_node): + left_node = updated_node.left + args = list(left_node.slice) + args[-1] = args[-1].with_changes( + comma=cst.Comma( + whitespace_before=cst.SimpleWhitespace( + value="", + ), + whitespace_after=cst.SimpleWhitespace( + value=" ", + ), + ) + ) + args.append( + cst.SubscriptElement( + slice=cst.Index(value=cst.Name(value="None")), + comma=cst.MaybeSentinel.DEFAULT, + ) + ) + left_node = left_node.with_changes(slice=tuple(args)) + return left_node diff --git a/tests/recorded/flatten_literal.txt b/tests/recorded/flatten_literal.txt new file mode 100644 index 0000000..0ca3534 --- /dev/null +++ b/tests/recorded/flatten_literal.txt @@ -0,0 +1,7 @@ +Literal[1, 2] | None # this should not change +var: Literal[1, 2] | None + +================================================================================ + +Literal[1, 2] | None # this should not change +var: Literal[1, 2, None] From ee6f6c1d0d17a6b202eb0a6cb87be0f5f503cc22 Mon Sep 17 00:00:00 2001 From: Cheukting Date: Wed, 11 May 2022 12:51:20 +0100 Subject: [PATCH 2/8] flatten union --- src/shed/_codemods.py | 22 ++++++++++++++++++++++ tests/recorded/flatten_union.txt | 5 +++++ 2 files changed, 27 insertions(+) create mode 100644 tests/recorded/flatten_union.txt diff --git a/src/shed/_codemods.py b/src/shed/_codemods.py index 192771e..c34b71b 100644 --- a/src/shed/_codemods.py +++ b/src/shed/_codemods.py @@ -362,3 +362,25 @@ def flatten_literal(self, _, updated_node): ) left_node = left_node.with_changes(slice=tuple(args)) return left_node + + @m.leave(m.Subscript(value=m.Name(value="Union"))) + def flatten_union(self, _, updated_node): + new_slice = [] + has_none = False + for item in updated_node.slice: + if m.matches(item.slice.value, m.Subscript(m.Name("Optional"))): + new_slice.append( + item.with_changes(slice=item.slice.value.slice[0].slice) + ) # peel off "Optional" + has_none = True + else: + new_slice.append(item) + if has_none: + new_slice.append( + cst.SubscriptElement( + slice=cst.Index(value=cst.Name(value="None")), + comma=cst.MaybeSentinel.DEFAULT, + ) + ) + return updated_node.with_changes(slice=new_slice) + return updated_node # nothing changes diff --git a/tests/recorded/flatten_union.txt b/tests/recorded/flatten_union.txt new file mode 100644 index 0000000..3c7a64b --- /dev/null +++ b/tests/recorded/flatten_union.txt @@ -0,0 +1,5 @@ +Union[int, Optional[str], bool] + +================================================================================ + +Union[int, str, bool, None] From b7ca7608cfea64dd5ae9d8ee276b32b8fb172942 Mon Sep 17 00:00:00 2001 From: Cheukting Date: Wed, 11 May 2022 23:45:07 +0100 Subject: [PATCH 3/8] refactor code, add test case --- src/shed/_codemods.py | 42 +++++++++++++----------------- tests/recorded/flatten_literal.txt | 4 +++ tests/recorded/flatten_union.txt | 4 +++ 3 files changed, 26 insertions(+), 24 deletions(-) diff --git a/src/shed/_codemods.py b/src/shed/_codemods.py index c34b71b..399b420 100644 --- a/src/shed/_codemods.py +++ b/src/shed/_codemods.py @@ -341,37 +341,33 @@ def _has_none(node): right=m.Name("None"), ) ) - def flatten_literal(self, _, updated_node): - left_node = updated_node.left - args = list(left_node.slice) - args[-1] = args[-1].with_changes( - comma=cst.Comma( - whitespace_before=cst.SimpleWhitespace( - value="", - ), - whitespace_after=cst.SimpleWhitespace( - value=" ", - ), - ) - ) + def flatten_literal_op(self, _, updated_node): + literal = updated_node.left + args = list(literal.slice) + for item in args: + if m.matches(item, m.SubscriptElement(m.Index(m.Name("None")))): + return literal # Already has "None" + args[-1] = args[-1].with_changes(comma=cst.Comma()) args.append( cst.SubscriptElement( slice=cst.Index(value=cst.Name(value="None")), - comma=cst.MaybeSentinel.DEFAULT, ) ) - left_node = left_node.with_changes(slice=tuple(args)) - return left_node + return literal.with_changes(slice=tuple(args)) - @m.leave(m.Subscript(value=m.Name(value="Union"))) - def flatten_union(self, _, updated_node): + @m.leave(m.Subscript(value=m.Name(value="Union") | m.Name(value="Literal"))) + def flatten_union_literal_subscript(self, _, updated_node): new_slice = [] has_none = False for item in updated_node.slice: if m.matches(item.slice.value, m.Subscript(m.Name("Optional"))): - new_slice.append( - item.with_changes(slice=item.slice.value.slice[0].slice) - ) # peel off "Optional" + new_slice += item.slice.value.slice # peel off "Optional" + has_none = True + elif m.matches( + item.slice.value, m.Subscript(m.Name("Union") | m.Name("Literal")) + ): + new_slice += item.slice.value.slice # peel off "Union" or "Literal" + elif m.matches(item.slice.value, m.Name("None")): has_none = True else: new_slice.append(item) @@ -379,8 +375,6 @@ def flatten_union(self, _, updated_node): new_slice.append( cst.SubscriptElement( slice=cst.Index(value=cst.Name(value="None")), - comma=cst.MaybeSentinel.DEFAULT, ) ) - return updated_node.with_changes(slice=new_slice) - return updated_node # nothing changes + return updated_node.with_changes(slice=new_slice) diff --git a/tests/recorded/flatten_literal.txt b/tests/recorded/flatten_literal.txt index 0ca3534..2057638 100644 --- a/tests/recorded/flatten_literal.txt +++ b/tests/recorded/flatten_literal.txt @@ -1,7 +1,11 @@ Literal[1, 2] | None # this should not change var: Literal[1, 2] | None +var2: Literal[1, Literal[2, 3]] +var3: Literal[None, 1, 2] | None ================================================================================ Literal[1, 2] | None # this should not change var: Literal[1, 2, None] +var2: Literal[1, 2, 3] +var3: Literal[1, 2, None] diff --git a/tests/recorded/flatten_union.txt b/tests/recorded/flatten_union.txt index 3c7a64b..3f601b8 100644 --- a/tests/recorded/flatten_union.txt +++ b/tests/recorded/flatten_union.txt @@ -1,5 +1,9 @@ Union[int, Optional[str], bool] +Union[None, int, Optional[str]] +Union[Union[int, float], str] ================================================================================ Union[int, str, bool, None] +Union[int, str, None] +Union[int, float, str] From 082ce78045f274325af09100086cb9a2bf0d36d5 Mon Sep 17 00:00:00 2001 From: Cheukting Date: Thu, 12 May 2022 12:28:00 +0100 Subject: [PATCH 4/8] adding more test --- src/shed/_codemods.py | 2 +- tests/recorded/flatten_literal.txt | 7 ++++++- tests/recorded/flatten_union.txt | 6 +++++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/src/shed/_codemods.py b/src/shed/_codemods.py index 399b420..2eca1b0 100644 --- a/src/shed/_codemods.py +++ b/src/shed/_codemods.py @@ -365,7 +365,7 @@ def flatten_union_literal_subscript(self, _, updated_node): has_none = True elif m.matches( item.slice.value, m.Subscript(m.Name("Union") | m.Name("Literal")) - ): + ) and updated_node.value == item.slice.value: new_slice += item.slice.value.slice # peel off "Union" or "Literal" elif m.matches(item.slice.value, m.Name("None")): has_none = True diff --git a/tests/recorded/flatten_literal.txt b/tests/recorded/flatten_literal.txt index 2057638..690f128 100644 --- a/tests/recorded/flatten_literal.txt +++ b/tests/recorded/flatten_literal.txt @@ -2,10 +2,15 @@ Literal[1, 2] | None # this should not change var: Literal[1, 2] | None var2: Literal[1, Literal[2, 3]] var3: Literal[None, 1, 2] | None +Literal[1, 2, Union[bool, str]] # this should not change +Literal[1, 2, Union[bool, str], Optional[int]] + ================================================================================ Literal[1, 2] | None # this should not change var: Literal[1, 2, None] -var2: Literal[1, 2, 3] +var2: Literal[1, Literal[2, 3]] var3: Literal[1, 2, None] +Literal[1, 2, Union[bool, str]] # this should not change +Literal[1, 2, Union[bool, str], int, None] diff --git a/tests/recorded/flatten_union.txt b/tests/recorded/flatten_union.txt index 3f601b8..c819c95 100644 --- a/tests/recorded/flatten_union.txt +++ b/tests/recorded/flatten_union.txt @@ -1,9 +1,13 @@ Union[int, Optional[str], bool] Union[None, int, Optional[str]] Union[Union[int, float], str] +Union[int, Literal[1, 2]] # this should not change +Union[int, Literal[1, 2], Optional[str]] ================================================================================ Union[int, str, bool, None] Union[int, str, None] -Union[int, float, str] +Union[Union[int, float], str] +Union[int, Literal[1, 2]] # this should not change +Union[int, Literal[1, 2], str, None] From f87f6e0d03e30255551457b835d1360f7984b9c2 Mon Sep 17 00:00:00 2001 From: Cheukting Date: Thu, 12 May 2022 12:35:26 +0100 Subject: [PATCH 5/8] remoce comma --- src/shed/_codemods.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/shed/_codemods.py b/src/shed/_codemods.py index 2eca1b0..f6b61d9 100644 --- a/src/shed/_codemods.py +++ b/src/shed/_codemods.py @@ -347,7 +347,6 @@ def flatten_literal_op(self, _, updated_node): for item in args: if m.matches(item, m.SubscriptElement(m.Index(m.Name("None")))): return literal # Already has "None" - args[-1] = args[-1].with_changes(comma=cst.Comma()) args.append( cst.SubscriptElement( slice=cst.Index(value=cst.Name(value="None")), From b72bde84e216bb1709401624060af80b53438c4b Mon Sep 17 00:00:00 2001 From: Cheukting Date: Thu, 12 May 2022 12:52:52 +0100 Subject: [PATCH 6/8] remove more comma --- src/shed/_codemods.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/shed/_codemods.py b/src/shed/_codemods.py index f6b61d9..8b9658a 100644 --- a/src/shed/_codemods.py +++ b/src/shed/_codemods.py @@ -130,8 +130,7 @@ def convert_optional_literal_to_literal_none(self, _, updated_node): args[-1] = args[-1].with_changes(comma=cst.Comma()) args.append( cst.SubscriptElement( - slice=cst.Index(value=cst.Name(value="None")), - comma=cst.MaybeSentinel.DEFAULT, + slice=cst.Index(value=cst.Name(value="None")) ) ) expr = expr.with_changes(slice=tuple(args)) @@ -305,9 +304,6 @@ def replace_unnecessary_listcomp_or_setcomp(self, _, updated_node): def reorder_union_literal_contents_none_last(self, _, updated_node): subscript_slice = list(updated_node.slice) subscript_slice.sort(key=lambda elt: elt.slice.value.value == "None") - subscript_slice[-1] = subscript_slice[-1].with_changes( - comma=cst.MaybeSentinel.DEFAULT - ) return updated_node.with_changes(slice=subscript_slice) @m.call_if_inside(m.Annotation(annotation=m.BinaryOperation())) From daa35b846aebba6ea8b2ccd85711d5a29b30fbde Mon Sep 17 00:00:00 2001 From: Cheukting Date: Thu, 12 May 2022 13:02:50 +0100 Subject: [PATCH 7/8] fix nexted Literal --- src/shed/_codemods.py | 2 +- tests/recorded/flatten_literal.txt | 3 +-- tests/recorded/flatten_union.txt | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/shed/_codemods.py b/src/shed/_codemods.py index 8b9658a..7d2198b 100644 --- a/src/shed/_codemods.py +++ b/src/shed/_codemods.py @@ -360,7 +360,7 @@ def flatten_union_literal_subscript(self, _, updated_node): has_none = True elif m.matches( item.slice.value, m.Subscript(m.Name("Union") | m.Name("Literal")) - ) and updated_node.value == item.slice.value: + ) and m.matches(updated_node.value, item.slice.value.value): new_slice += item.slice.value.slice # peel off "Union" or "Literal" elif m.matches(item.slice.value, m.Name("None")): has_none = True diff --git a/tests/recorded/flatten_literal.txt b/tests/recorded/flatten_literal.txt index 690f128..0a41277 100644 --- a/tests/recorded/flatten_literal.txt +++ b/tests/recorded/flatten_literal.txt @@ -5,12 +5,11 @@ var3: Literal[None, 1, 2] | None Literal[1, 2, Union[bool, str]] # this should not change Literal[1, 2, Union[bool, str], Optional[int]] - ================================================================================ Literal[1, 2] | None # this should not change var: Literal[1, 2, None] -var2: Literal[1, Literal[2, 3]] +var2: Literal[1, 2, 3] var3: Literal[1, 2, None] Literal[1, 2, Union[bool, str]] # this should not change Literal[1, 2, Union[bool, str], int, None] diff --git a/tests/recorded/flatten_union.txt b/tests/recorded/flatten_union.txt index c819c95..bd21a0e 100644 --- a/tests/recorded/flatten_union.txt +++ b/tests/recorded/flatten_union.txt @@ -8,6 +8,6 @@ Union[int, Literal[1, 2], Optional[str]] Union[int, str, bool, None] Union[int, str, None] -Union[Union[int, float], str] +Union[int, float, str] Union[int, Literal[1, 2]] # this should not change Union[int, Literal[1, 2], str, None] From 551b94f2b08038344e2fc8afeb525aad615d824d Mon Sep 17 00:00:00 2001 From: Cheukting Date: Thu, 12 May 2022 13:04:54 +0100 Subject: [PATCH 8/8] shed my beautiful stuff --- src/shed/_codemods.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/shed/_codemods.py b/src/shed/_codemods.py index 7d2198b..09cd198 100644 --- a/src/shed/_codemods.py +++ b/src/shed/_codemods.py @@ -128,11 +128,7 @@ def convert_optional_literal_to_literal_none(self, _, updated_node): expr = updated_node.slice[0].slice.value args = list(expr.slice) args[-1] = args[-1].with_changes(comma=cst.Comma()) - args.append( - cst.SubscriptElement( - slice=cst.Index(value=cst.Name(value="None")) - ) - ) + args.append(cst.SubscriptElement(slice=cst.Index(value=cst.Name(value="None")))) expr = expr.with_changes(slice=tuple(args)) return expr