From 183059f9a46b17bb8bd3a7fb4b804b6255c0d2b4 Mon Sep 17 00:00:00 2001
From: Michael Carlstrom <rmc@carlstrom.com>
Date: Tue, 25 Jun 2024 17:57:34 -0400
Subject: [PATCH] feat(types): add support for typing.Literal type (#5192)

* typevar prototype

* style: pre-commit fixes

* change to NameT

* style: pre-commit fixes

* make string const

* add missing closing bracket

* style: pre-commit fixes

* clean up handle_type_name

* style: pre-commit fixes

* add back missing <

* style: pre-commit fixes

* add back NameT

* try fixed_string

* style: pre-commit fixes

* std::basic_fixed_string

* test c++20

* style: pre-commit fixes

* cleanup

* fix object to typevar conversion

* style: pre-commit fixes

* And CPP20 checks

* style: pre-commit fixes

* add missing cpp20++ check

* style: pre-commit fixes

* Add C++20 check to python

* Fix python if {

* style: pre-commit fixes

* update test name

* style: pre-commit fixes

* remove call on cpp_std

* make field const

* test nontype_template

* update feature check

* update name of guard

* fix try except in test

* fix pre commit

* remove extra semi colon

* except AttributeError

* fix try except in test

* remove const

* Clean up tests

* style: pre-commit fixes

* start string literal

* start int literal

* func declare

* commit clean

* use contextlib.suppres

* resolve stash

* more literal type

* fix annotation name

* stash

* request changes

* lint

* Add comments

* style: pre-commit fixes

* Add support for unions and optionals to be compatible with object

* lint

* remove comment

* Create Literal Type implementation

* clean up

* Update comment

* remove incorrect comment

* rerun CI

* rerun CI

* fix extra line

* lint

* move if defined block

* style: pre-commit fixes

---------

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

diff --git a/include/pybind11/typing.h b/include/pybind11/typing.h
index 8d91c51d9..fb37f0021 100644
--- a/include/pybind11/typing.h
+++ b/include/pybind11/typing.h
@@ -83,8 +83,13 @@ class Optional : public object {
 #if defined(__cpp_nontype_template_parameter_class)
 template <size_t N>
 struct StringLiteral {
-    constexpr StringLiteral(const char (&str)[N]) { std::copy_n(str, N, value); }
-    char value[N];
+    constexpr StringLiteral(const char (&str)[N]) { std::copy_n(str, N, name); }
+    char name[N];
+};
+
+template <StringLiteral... StrLits>
+class Literal : public object {
+    PYBIND11_OBJECT_DEFAULT(Literal, object, PyObject_Type)
 };
 
 // Example syntax for creating a TypeVar.
@@ -172,9 +177,15 @@ struct handle_type_name<typing::Optional<T>> {
 };
 
 #if defined(__cpp_nontype_template_parameter_class)
+template <typing::StringLiteral... Literals>
+struct handle_type_name<typing::Literal<Literals...>> {
+    static constexpr auto name = const_name("Literal[")
+                                 + pybind11::detail::concat(const_name(Literals.name)...)
+                                 + const_name("]");
+};
 template <typing::StringLiteral StrLit>
 struct handle_type_name<typing::TypeVar<StrLit>> {
-    static constexpr auto name = const_name(StrLit.value);
+    static constexpr auto name = const_name(StrLit.name);
 };
 #endif
 
diff --git a/tests/test_pytypes.cpp b/tests/test_pytypes.cpp
index ce003e092..952b5af6e 100644
--- a/tests/test_pytypes.cpp
+++ b/tests/test_pytypes.cpp
@@ -110,6 +110,19 @@ void m_defs(py::module_ &m) {
 } // namespace handle_from_move_only_type_with_operator_PyObject
 
 #if defined(__cpp_nontype_template_parameter_class)
+namespace literals {
+enum Color { RED = 0, BLUE = 1 };
+
+typedef py::typing::Literal<"26",
+                            "0x1A",
+                            "\"hello world\"",
+                            "b\"hello world\"",
+                            "u\"hello world\"",
+                            "True",
+                            "Color.RED",
+                            "None">
+    LiteralFoo;
+} // namespace literals
 namespace typevar {
 typedef py::typing::TypeVar<"T"> TypeVarT;
 typedef py::typing::TypeVar<"V"> TypeVarV;
@@ -851,6 +864,7 @@ 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_type", [](const py::typing::Type<int> &t) -> py::type { return t; });
 
     m.def("annotate_union",
@@ -881,6 +895,11 @@ TEST_SUBMODULE(pytypes, m) {
           [](py::typing::Optional<int> &o) -> py::object { return o; });
 
 #if defined(__cpp_nontype_template_parameter_class)
+    py::enum_<literals::Color>(m, "Color")
+        .value("RED", literals::Color::RED)
+        .value("BLUE", literals::Color::BLUE);
+
+    m.def("annotate_literal", [](literals::LiteralFoo &o) -> py::object { return o; });
     m.def("annotate_generic_containers",
           [](const py::typing::List<typevar::TypeVarT> &l) -> py::typing::List<typevar::TypeVarV> {
               return l;
diff --git a/tests/test_pytypes.py b/tests/test_pytypes.py
index 19e002de9..b265512c8 100644
--- a/tests/test_pytypes.py
+++ b/tests/test_pytypes.py
@@ -998,6 +998,17 @@ def test_optional_object_annotations(doc):
     )
 
 
+@pytest.mark.skipif(
+    not m.if_defined__cpp_nontype_template_parameter_class,
+    reason="C++20 feature not available.",
+)
+def test_literal(doc):
+    assert (
+        doc(m.annotate_literal)
+        == 'annotate_literal(arg0: Literal[26, 0x1A, "hello world", b"hello world", u"hello world", True, Color.RED, None]) -> object'
+    )
+
+
 @pytest.mark.skipif(
     not m.if_defined__cpp_nontype_template_parameter_class,
     reason="C++20 feature not available.",