From 9519e3652b7aa4b9b5c8c640fca5138744f03b88 Mon Sep 17 00:00:00 2001
From: Graham Bleaney <gbleaney@fb.com>
Date: Mon, 18 Jul 2022 22:34:27 -0400
Subject: [PATCH] Add LiteralString support to string module (#8268)

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
---
 stdlib/_typeshed/__init__.pyi |  4 ++++
 stdlib/string.pyi             | 36 +++++++++++++++++++++++------------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/stdlib/_typeshed/__init__.pyi b/stdlib/_typeshed/__init__.pyi
index ad78640b4..89ca9d816 100644
--- a/stdlib/_typeshed/__init__.pyi
+++ b/stdlib/_typeshed/__init__.pyi
@@ -266,6 +266,10 @@ class structseq(Generic[_T_co]):
 # Superset of typing.AnyStr that also inclues LiteralString
 AnyOrLiteralStr = TypeVar("AnyOrLiteralStr", str, bytes, LiteralString)  # noqa: Y001
 
+# Represents when str or LiteralStr is acceptable. Useful for string processing
+# APIs where literalness of return value depends on literalness of inputs
+StrOrLiteralStr = TypeVar("StrOrLiteralStr", LiteralString, str)  # noqa: Y001
+
 # Objects suitable to be passed to sys.setprofile, threading.setprofile, and similar
 ProfileFunction: TypeAlias = Callable[[FrameType, str, Any], object]
 
diff --git a/stdlib/string.pyi b/stdlib/string.pyi
index 320178f44..222129c7f 100644
--- a/stdlib/string.pyi
+++ b/stdlib/string.pyi
@@ -1,7 +1,9 @@
 import sys
+from _typeshed import StrOrLiteralStr
 from collections.abc import Iterable, Mapping, Sequence
 from re import Pattern, RegexFlag
-from typing import Any
+from typing import Any, overload
+from typing_extensions import LiteralString
 
 __all__ = [
     "ascii_letters",
@@ -18,17 +20,17 @@ __all__ = [
     "Template",
 ]
 
-ascii_letters: str
-ascii_lowercase: str
-ascii_uppercase: str
-digits: str
-hexdigits: str
-octdigits: str
-punctuation: str
-printable: str
-whitespace: str
+ascii_letters: LiteralString
+ascii_lowercase: LiteralString
+ascii_uppercase: LiteralString
+digits: LiteralString
+hexdigits: LiteralString
+octdigits: LiteralString
+punctuation: LiteralString
+printable: LiteralString
+whitespace: LiteralString
 
-def capwords(s: str, sep: str | None = ...) -> str: ...
+def capwords(s: StrOrLiteralStr, sep: StrOrLiteralStr | None = ...) -> StrOrLiteralStr: ...
 
 class Template:
     template: str
@@ -46,9 +48,19 @@ class Template:
 
 # TODO(MichalPokorny): This is probably badly and/or loosely typed.
 class Formatter:
+    @overload
+    def format(self, __format_string: LiteralString, *args: LiteralString, **kwargs: LiteralString) -> LiteralString: ...
+    @overload
     def format(self, __format_string: str, *args: Any, **kwargs: Any) -> str: ...
+    @overload
+    def vformat(
+        self, format_string: LiteralString, args: Sequence[LiteralString], kwargs: Mapping[LiteralString, LiteralString]
+    ) -> LiteralString: ...
+    @overload
     def vformat(self, format_string: str, args: Sequence[Any], kwargs: Mapping[str, Any]) -> str: ...
-    def parse(self, format_string: str) -> Iterable[tuple[str, str | None, str | None, str | None]]: ...
+    def parse(
+        self, format_string: StrOrLiteralStr
+    ) -> Iterable[tuple[StrOrLiteralStr, StrOrLiteralStr | None, StrOrLiteralStr | None, StrOrLiteralStr | None]]: ...
     def get_field(self, field_name: str, args: Sequence[Any], kwargs: Mapping[str, Any]) -> Any: ...
     def get_value(self, key: int | str, args: Sequence[Any], kwargs: Mapping[str, Any]) -> Any: ...
     def check_unused_args(self, used_args: Sequence[int | str], args: Sequence[Any], kwargs: Mapping[str, Any]) -> None: ...