From 68405a1141fba1314c378133f12808a4d97bd3d8 Mon Sep 17 00:00:00 2001
From: Michael Carlstrom <rmc@carlstrom.com>
Date: Sat, 15 Jun 2024 00:16:26 -0400
Subject: [PATCH] Add `Union` and `Optional` to typing.h (#5165)

* add type unions and optionals

* add type inheritance

* style: pre-commit fixes

* switch to inheriting from object

* style: pre-commit fixes

* fix text case

* style: pre-commit fixes

* fix bind call

* fix function name

* add std::move for older code

* remove std::move

* move away from object creation

* style: pre-commit fixes

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 include/pybind11/typing.h | 22 ++++++++++++++++++++++
 tests/test_pytypes.cpp    | 22 ++++++++++++++++++++++
 tests/test_pytypes.py     | 21 +++++++++++++++++++++
 3 files changed, 65 insertions(+)

diff --git a/include/pybind11/typing.h b/include/pybind11/typing.h
index bc275fc5..2c9eaffb 100644
--- a/include/pybind11/typing.h
+++ b/include/pybind11/typing.h
@@ -63,6 +63,16 @@ class Callable<Return(Args...)> : public function {
     using function::function;
 };
 
+template <typename... Types>
+class Union : public object {
+    using object::object;
+};
+
+template <typename T>
+class Optional : public object {
+    using object::object;
+};
+
 PYBIND11_NAMESPACE_END(typing)
 
 PYBIND11_NAMESPACE_BEGIN(detail)
@@ -121,5 +131,17 @@ struct handle_type_name<typing::Callable<Return(Args...)>> {
           + const_name("], ") + make_caster<retval_type>::name + const_name("]");
 };
 
+template <typename... Types>
+struct handle_type_name<typing::Union<Types...>> {
+    static constexpr auto name = const_name("Union[")
+                                 + ::pybind11::detail::concat(make_caster<Types>::name...)
+                                 + const_name("]");
+};
+
+template <typename T>
+struct handle_type_name<typing::Optional<T>> {
+    static constexpr auto name = const_name("Optional[") + make_caster<T>::name + const_name("]");
+};
+
 PYBIND11_NAMESPACE_END(detail)
 PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)
diff --git a/tests/test_pytypes.cpp b/tests/test_pytypes.cpp
index f3709d40..e97dceee 100644
--- a/tests/test_pytypes.cpp
+++ b/tests/test_pytypes.cpp
@@ -844,4 +844,26 @@ TEST_SUBMODULE(pytypes, m) {
     m.def("annotate_iterator_int", [](const py::typing::Iterator<int> &) {});
     m.def("annotate_fn",
           [](const py::typing::Callable<int(py::typing::List<py::str>, py::str)> &) {});
+
+    m.def("annotate_union",
+          [](py::typing::List<py::typing::Union<py::str, py::int_, py::object>> l,
+             py::str a,
+             py::int_ b,
+             py::object c) -> py::typing::List<py::typing::Union<py::str, py::int_, py::object>> {
+              l.append(a);
+              l.append(b);
+              l.append(c);
+              return l;
+          });
+
+    m.def("union_typing_only",
+          [](py::typing::List<py::typing::Union<py::str>> &l)
+              -> py::typing::List<py::typing::Union<py::int_>> { return l; });
+
+    m.def("annotate_optional",
+          [](py::list &list) -> py::typing::List<py::typing::Optional<py::str>> {
+              list.append(py::str("hi"));
+              list.append(py::none());
+              return list;
+          });
 }
diff --git a/tests/test_pytypes.py b/tests/test_pytypes.py
index 38edfd99..ca2c51c4 100644
--- a/tests/test_pytypes.py
+++ b/tests/test_pytypes.py
@@ -955,3 +955,24 @@ def test_fn_annotations(doc):
         doc(m.annotate_fn)
         == "annotate_fn(arg0: Callable[[list[str], str], int]) -> None"
     )
+
+
+def test_union_annotations(doc):
+    assert (
+        doc(m.annotate_union)
+        == "annotate_union(arg0: list[Union[str, int, object]], arg1: str, arg2: int, arg3: object) -> list[Union[str, int, object]]"
+    )
+
+
+def test_union_typing_only(doc):
+    assert (
+        doc(m.union_typing_only)
+        == "union_typing_only(arg0: list[Union[str]]) -> list[Union[int]]"
+    )
+
+
+def test_optional_annotations(doc):
+    assert (
+        doc(m.annotate_optional)
+        == "annotate_optional(arg0: list) -> list[Optional[str]]"
+    )