From ff19629b119f63d35c76d3d7e8dba5ea2db11ab7 Mon Sep 17 00:00:00 2001
From: Shaygan Hooshyari <sh.hooshyari@gmail.com>
Date: Sun, 17 Nov 2024 18:04:58 +0100
Subject: [PATCH] Understand `typing.Optional` in annotations (#14397)

---
 .../resources/mdtest/annotations/optional.md  | 47 +++++++++++++++++++
 crates/red_knot_python_semantic/src/types.rs  |  9 +++-
 .../src/types/infer.rs                        |  4 ++
 .../red_knot_python_semantic/src/types/mro.rs |  5 +-
 4 files changed, 61 insertions(+), 4 deletions(-)
 create mode 100644 crates/red_knot_python_semantic/resources/mdtest/annotations/optional.md

diff --git a/crates/red_knot_python_semantic/resources/mdtest/annotations/optional.md b/crates/red_knot_python_semantic/resources/mdtest/annotations/optional.md
new file mode 100644
index 0000000000..f4a6962e7f
--- /dev/null
+++ b/crates/red_knot_python_semantic/resources/mdtest/annotations/optional.md
@@ -0,0 +1,47 @@
+# Optional
+
+## Annotation
+
+`typing.Optional` is equivalent to using the type with a None in a Union.
+
+```py
+from typing import Optional
+
+a: Optional[int]
+a1: Optional[bool]
+a2: Optional[Optional[bool]]
+a3: Optional[None]
+
+def f():
+    # revealed: int | None
+    reveal_type(a)
+    # revealed: bool | None
+    reveal_type(a1)
+    # revealed: bool | None
+    reveal_type(a2)
+    # revealed: None
+    reveal_type(a3)
+```
+
+## Assignment
+
+```py
+from typing import Optional
+
+a: Optional[int] = 1
+a = None
+# error: [invalid-assignment] "Object of type `Literal[""]` is not assignable to `int | None`"
+a = ""
+```
+
+## Typing Extensions
+
+```py
+from typing_extensions import Optional
+
+a: Optional[int]
+
+def f():
+    # revealed: int | None
+    reveal_type(a)
+```
diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs
index dc0f5196fb..a1a1f32d1e 100644
--- a/crates/red_knot_python_semantic/src/types.rs
+++ b/crates/red_knot_python_semantic/src/types.rs
@@ -1807,6 +1807,8 @@ impl<'db> KnownClass {
 pub enum KnownInstanceType<'db> {
     /// The symbol `typing.Literal` (which can also be found as `typing_extensions.Literal`)
     Literal,
+    /// The symbol `typing.Optional` (which can also be found as `typing_extensions.Literal`)
+    Optional,
     /// A single instance of `typing.TypeVar`
     TypeVar(TypeVarInstance<'db>),
     // TODO: fill this enum out with more special forms, etc.
@@ -1816,6 +1818,7 @@ impl<'db> KnownInstanceType<'db> {
     pub const fn as_str(self) -> &'static str {
         match self {
             KnownInstanceType::Literal => "Literal",
+            KnownInstanceType::Optional => "Optional",
             KnownInstanceType::TypeVar(_) => "TypeVar",
         }
     }
@@ -1823,8 +1826,7 @@ impl<'db> KnownInstanceType<'db> {
     /// Evaluate the known instance in boolean context
     pub const fn bool(self) -> Truthiness {
         match self {
-            Self::Literal => Truthiness::AlwaysTrue,
-            Self::TypeVar(_) => Truthiness::AlwaysTrue,
+            Self::Literal | Self::Optional | Self::TypeVar(_) => Truthiness::AlwaysTrue,
         }
     }
 
@@ -1832,6 +1834,7 @@ impl<'db> KnownInstanceType<'db> {
     pub fn repr(self, db: &'db dyn Db) -> &'db str {
         match self {
             Self::Literal => "typing.Literal",
+            Self::Optional => "typing.Optional",
             Self::TypeVar(typevar) => typevar.name(db),
         }
     }
@@ -1840,6 +1843,7 @@ impl<'db> KnownInstanceType<'db> {
     pub const fn class(self) -> KnownClass {
         match self {
             Self::Literal => KnownClass::SpecialForm,
+            Self::Optional => KnownClass::SpecialForm,
             Self::TypeVar(_) => KnownClass::TypeVar,
         }
     }
@@ -1859,6 +1863,7 @@ impl<'db> KnownInstanceType<'db> {
         }
         match (module.name().as_str(), instance_name) {
             ("typing" | "typing_extensions", "Literal") => Some(Self::Literal),
+            ("typing" | "typing_extensions", "Optional") => Some(Self::Optional),
             _ => None,
         }
     }
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
index bc51a25be7..59d3fa8023 100644
--- a/crates/red_knot_python_semantic/src/types/infer.rs
+++ b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -4524,6 +4524,10 @@ impl<'db> TypeInferenceBuilder<'db> {
                     Type::Unknown
                 }
             },
+            KnownInstanceType::Optional => {
+                let param_type = self.infer_type_expression(parameters);
+                UnionType::from_elements(self.db, [param_type, Type::none(self.db)])
+            }
             KnownInstanceType::TypeVar(_) => Type::Todo,
         }
     }
diff --git a/crates/red_knot_python_semantic/src/types/mro.rs b/crates/red_knot_python_semantic/src/types/mro.rs
index 68e092d9d8..7f5282ce9e 100644
--- a/crates/red_knot_python_semantic/src/types/mro.rs
+++ b/crates/red_knot_python_semantic/src/types/mro.rs
@@ -371,8 +371,9 @@ impl<'db> ClassBase<'db> {
             | Type::ModuleLiteral(_)
             | Type::SubclassOf(_) => None,
             Type::KnownInstance(known_instance) => match known_instance {
-                KnownInstanceType::Literal => None,
-                KnownInstanceType::TypeVar(_) => None,
+                KnownInstanceType::TypeVar(_)
+                | KnownInstanceType::Literal
+                | KnownInstanceType::Optional => None,
             },
         }
     }