From cfd7268d7f7961ffdd1948174763d3c2ad78a793 Mon Sep 17 00:00:00 2001
From: JoshuaTianYangLiu <josh.liu@student.tdsb.on.ca>
Date: Mon, 2 May 2022 17:29:38 -0600
Subject: [PATCH] Rewrite `Optional[Literal[...]]` to `Literal[..., None]` #22

---
 src/shed/_codemods.py                | 25 +++++++++++++++++++++++++
 tests/recorded/optionals_literal.txt |  7 +++++++
 2 files changed, 32 insertions(+)
 create mode 100644 tests/recorded/optionals_literal.txt

diff --git a/src/shed/_codemods.py b/src/shed/_codemods.py
index 0b2fc46..d76962c 100644
--- a/src/shed/_codemods.py
+++ b/src/shed/_codemods.py
@@ -56,6 +56,7 @@ def _run_codemods(code: str, min_version: Tuple[int, int]) -> str:
 
     if imports_hypothesis(code):  # pragma: no cover
         mod = attempt_hypothesis_codemods(context, mod)
+    # print(mod)
     mod = ShedFixers(context).transform_module(mod)
     return mod.code
 
@@ -114,6 +115,30 @@ def leave_Assert(self, _, updated_node):  # noqa
             cst.Call(cst.Name("AssertionError"), args=[cst.Arg(updated_node.msg)])
         )
 
+    @m.leave(
+        m.Subscript(
+            value=m.Name(value="Optional"),
+            slice=[
+                m.SubscriptElement(
+                    slice=m.Index(value=m.Subscript(value=m.Name(value="Literal")))
+                )
+            ],
+        )
+    )
+    def convert_optional_literal_to_literal_none(self, _, updated_node):
+        """Inspired by Pybetter."""
+        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")),
+                comma=cst.MaybeSentinel.DEFAULT,
+            )
+        )
+        expr = expr.with_changes(slice=tuple(args))
+        return expr
+
     @m.leave(m.ComparisonTarget(comparator=m.Name(value="None"), operator=m.Equal()))
     def convert_none_cmp(self, _, updated_node):
         """Inspired by Pybetter."""
diff --git a/tests/recorded/optionals_literal.txt b/tests/recorded/optionals_literal.txt
new file mode 100644
index 0000000..81c7bc5
--- /dev/null
+++ b/tests/recorded/optionals_literal.txt
@@ -0,0 +1,7 @@
+Optional[Literal[1, 2]]  # Literal[1, 2, None]
+var: Optional[Literal[1, 2]] = 1  # Literal[1, 2, None]
+
+================================================================================
+
+Literal[1, 2, None]  # Literal[1, 2, None]
+var: Literal[1, 2, None] = 1  # Literal[1, 2, None]