Improve type inference for f-strings containing literals (#571)

Author: JelleZijlstra

docs/changelog.md

@@ -2,7 +2,8 @@
 
 ## Unreleased
 
-- Add experimental `@has_extra_keys` decorator for `TypedDictt` types (#568)
+- Improve type inference for f-strings containing literals (#571)
+- Add experimental `@has_extra_keys` decorator for `TypedDict` types (#568)
 - Fix crash on recursive type aliases. Recursive type aliases now fall back to `Any` (#565)
 - Support `in` on objects with only `__getitem__` (#564)
 - Add support for `except*` (PEP 654) (#562)

pyanalyze/name_check_visitor.py

@@ -2579,11 +2579,77 @@ def _visit_comprehension_inner(
     # Literals and displays
 
     def visit_JoinedStr(self, node: ast.JoinedStr) -> Value:
-        # JoinedStr is the node type for f-strings.
-        # Not too much to check here. Perhaps we can add checks that format specifiers
-        # are valid.
-        self._generic_visit_list(node.values)
-        return TypedValue(str)
+        elements = self._generic_visit_list(node.values)
+        limit = self.options.get_value_for(UnionSimplificationLimit)
+        possible_values: List[List[str]] = [[]]
+        for elt in elements:
+            subvals = list(flatten_values(elt))
+            # Bail out if the list of possible values gets too long.
+            if len(possible_values) * len(subvals) > limit:
+                return TypedValue(str)
+            to_add = []
+            for subval in subvals:
+                if not isinstance(subval, KnownValue):
+                    return TypedValue(str)
+                if not isinstance(subval.val, str):
+                    return TypedValue(str)
+                to_add.append(subval.val)
+            possible_values = [
+                lst + [new_elt] for lst in possible_values for new_elt in to_add
+            ]
+        return unite_values(*[KnownValue("".join(lst)) for lst in possible_values])
+
+    def visit_FormattedValue(self, node: ast.FormattedValue) -> Value:
+        val = self.visit(node.value)
+        format_spec_val = (
+            self.visit(node.format_spec) if node.format_spec else KnownValue("")
+        )
+        if isinstance(format_spec_val, KnownValue) and isinstance(
+            format_spec_val.val, str
+        ):
+            format_spec = format_spec_val.val
+        else:
+            # TODO: statically check whether the format specifier is valid.
+            return TypedValue(str)
+        possible_vals = []
+        for subval in flatten_values(val):
+            possible_vals.append(
+                self._visit_single_formatted_value(subval, node, format_spec)
+            )
+        return unite_and_simplify(
+            *possible_vals, limit=self.options.get_value_for(UnionSimplificationLimit)
+        )
+
+    def _visit_single_formatted_value(
+        self, val: Value, node: ast.FormattedValue, format_spec: str
+    ) -> Value:
+        if not isinstance(val, KnownValue):
+            return TypedValue(str)
+        output = val.val
+        if node.conversion != -1:
+            unsupported_conversion = False
+            try:
+                if node.conversion == ord("a"):
+                    output = ascii(output)
+                elif node.conversion == ord("s"):
+                    output = str(output)
+                elif node.conversion == ord("r"):
+                    output = repr(output)
+                else:
+                    unsupported_conversion = True
+            except Exception:
+                # str/repr/ascii failed
+                return TypedValue(str)
+            if unsupported_conversion:
+                raise NotImplementedError(
+                    f"Unsupported converion specifier {node.conversion}"
+                )
+        try:
+            output = format(output, format_spec)
+        except Exception:
+            # format failed
+            return TypedValue(str)
+        return KnownValue(output)
 
     def visit_Constant(self, node: ast.Constant) -> Value:
         # replaces Num, Str, etc. in 3.8+

pyanalyze/test_format_strings.py

@@ -478,3 +478,30 @@ def capybara(x):
     def test_complicated_expression(self):
         def capybara(x):
             "foo %s" % len(x)
+
+
+class TestFStringLiteral(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_basic(self):
+        from typing_extensions import Literal, assert_type
+
+        def capybara(a: Literal["a"], ab: Literal["a", "b"]):
+            assert_type(f"a{a}", Literal["aa"])
+            assert_type(f"a{ab}", Literal["aa", "ab"])
+            assert_type(f"a{ab}b{ab}", Literal["aaba", "aabb", "abba", "abbb"])
+            # Make sure we don't infer a 2**16-size union
+            assert_type(
+                f"{ab}{ab}{ab}{ab}{ab}{ab}{ab}{ab}{ab}{ab}{ab}{ab}{ab}{ab}{ab}{ab}", str
+            )
+
+    @assert_passes()
+    def test_conversions(self):
+        from typing_extensions import Literal, assert_type
+
+        def capybara(a: Literal["á"], ab: Literal["á", "ê"]):
+            assert_type(f"a{a!r}", Literal["a'á'"])
+            assert_type(f"a{ab!r}", Literal["a'á'", "a'ê'"])
+            assert_type(f"a{a!s}", Literal["aá"])
+            assert_type(f"a{ab!s}", Literal["aá", "aê"])
+            assert_type(f"a{a!a}", Literal["a'\\xe1'"])
+            assert_type(f"a{ab!a}", Literal["a'\\xe1'", "a'\\xea'"])