Merge pull request #5 from VallinZ/AddAlignmentSpecifierLogic

Extend check on None for str.format()

Author: ChristinaTrinh

mypy/checkstrformat.py

@@ -337,6 +337,7 @@ def check_specs_in_format_call(
 
         The core logic for format checking is implemented in this method.
         """
+
         assert all(s.key for s in specs), "Keys must be auto-generated first!"
         replacements = self.find_replacements_in_call(call, [cast(str, s.key) for s in specs])
         assert len(replacements) == len(specs)
@@ -448,24 +449,27 @@ def perform_special_format_checks(
                     call,
                     code=codes.STRING_FORMATTING,
                 )
-        a_type = get_proper_type(actual_type)
-        if isinstance(a_type, NoneType):
+
+        if isinstance(get_proper_type(actual_type), NoneType):
             # Perform type check of alignment specifiers on None
-            if spec.format_spec and any(c in spec.format_spec for c in "<>^"):
-                specifierIndex = -1
-                for i in range(len("<>^")):
-                    if spec.format_spec[i] in "<>^":
-                        specifierIndex = i
-                if specifierIndex > -1:
-                    self.msg.fail(
-                        (
-                            f"Alignment format specifier "
-                            f'"{spec.format_spec[specifierIndex]}" '
-                            f"is not supported for None"
-                        ),
-                        call,
-                        code=codes.STRING_FORMATTING,
-                    )
+            # If spec.format_spec is None then we use "" instead of avoid crashing
+            specifier_char = None
+            if spec.non_standard_format_spec and isinstance(call.args[-1], StrExpr):
+                arg = call.args[-1].value
+                specifier_char = next((c for c in (arg or "") if c in "<>^"), None)
+            elif isinstance(spec.format_spec, str):
+                specifier_char = next((c for c in (spec.format_spec or "") if c in "<>^"), None)
+
+            if specifier_char:
+                self.msg.fail(
+                    (
+                        f"Alignment format specifier "
+                        f'"{specifier_char}" '
+                        f"is not supported for None"
+                    ),
+                    call,
+                    code=codes.STRING_FORMATTING,
+                )
 
     def find_replacements_in_call(self, call: CallExpr, keys: list[str]) -> list[Expression]:
         """Find replacement expression for every specifier in str.format() call.

test-data/unit/check-formatting.test

@@ -266,13 +266,34 @@ b'%c' % (123)
 
 
 [case testFormatCallNoneAlignment]
+from typing import Optional
+
 '{:<1}'.format(None)  # E: Alignment format specifier "<" is not supported for None
 '{:>1}'.format(None)  # E: Alignment format specifier ">" is not supported for None
 '{:^1}'.format(None)  # E: Alignment format specifier "^" is not supported for None
 
 '{:<10}'.format('16') # OK
 '{:>10}'.format('16') # OK
 '{:^10}'.format('16') # OK
+
+'{!s:<5}'.format(None) # OK
+'{!s:>5}'.format(None) # OK
+'{!s:^5}'.format(None) # OK
+
+f"{None!s:<5}"  # OK
+f"{None!s:>5}"  # OK
+f"{None!s:^5}"  # OK
+
+
+f"{None:<5}"  # E: Alignment format specifier "<" is not supported for None
+f"{None:>5}"  # E: Alignment format specifier ">" is not supported for None
+f"{None:^5}"  # E: Alignment format specifier "^" is not supported for None
+
+my_var: Optional[str] = None
+"{:<2}".format(my_var) # E: Alignment format specifier "<" is not supported for None
+my_var = "test"
+"{:>2}".format(my_var) # OK
+
 [builtins fixtures/primitives.pyi]
 
 [case testFormatCallParseErrors]