Flatten syntactically nested Union or Literal instances

src/shed/_codemods.py

@@ -332,3 +332,55 @@ 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
+
+    @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

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]

tests/recorded/flatten_union.txt

@@ -0,0 +1,5 @@
+Union[int, Optional[str], bool]
+
+================================================================================
+
+Union[int, str, bool, None]