Repository: quora/pyanalyze
Pull request #816 (2024-09-27 17:24:58)

Flag invalid regexes


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 82dfd08..4c89391 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -1,5 +1,10 @@
 # Changelog
 
+## Unreleased
+
+- Flag invalid regexes in arguments to functions like
+  `re.search` (#816)
+
 ## Version 0.13.1 (August 7, 2024)
 
 - Use Trusted Publishing to publish releases (#806)
diff --git a/pyanalyze/__init__.py b/pyanalyze/__init__.py
index f162194..edce733 100644
--- a/pyanalyze/__init__.py
+++ b/pyanalyze/__init__.py
@@ -58,6 +58,7 @@ used(extensions.is_of_type)
 used(extensions.show_error)
 used(extensions.has_extra_keys)
 used(extensions.EnumName)
+used(extensions.ValidRegex)
 used(value.UNRESOLVED_VALUE)  # keeping it around for now just in case
 used(reexport)
 used(patma)
diff --git a/pyanalyze/arg_spec.py b/pyanalyze/arg_spec.py
index 563adf3..92d857e 100644
--- a/pyanalyze/arg_spec.py
+++ b/pyanalyze/arg_spec.py
@@ -330,7 +330,7 @@ class UnwrapClass(PyObjectSequenceOption[_Unwrapper]):
         return typ
 
 
-_BUILTIN_KNOWN_SIGNATURES = []
+_BUILTIN_KNOWN_SIGNATURES = [implementation.get_default_argspecs_with_cache]
 
 try:
     import _pytest
@@ -784,6 +784,13 @@ class ArgSpecCache:
             obj, allow_call=allow_call, type_params=type_params
         )
         if argspec is not None:
+            if impl is not None:
+                if isinstance(argspec, OverloadedSignature):
+                    return OverloadedSignature(
+                        [replace(sig, impl=impl) for sig in argspec.signatures]
+                    )
+                else:
+                    return replace(argspec, impl=impl)
             return argspec
 
         if is_typeddict(obj) and not is_typing_name(obj, "TypedDict"):
diff --git a/pyanalyze/error_code.py b/pyanalyze/error_code.py
index f3d08aa..7d8a649 100644
--- a/pyanalyze/error_code.py
+++ b/pyanalyze/error_code.py
@@ -170,6 +170,7 @@ ErrorCode = ErrorRegistry(
         Error("generator_return", "Generator must return an iterable"),
         Error("unsafe_comparison", "Non-overlapping equality checks"),
         Error("must_use", "Value cannot be discarded"),
+        Error("invalid_regex", "Invalid regular expression"),
     ]
 )
 
diff --git a/pyanalyze/extensions.py b/pyanalyze/extensions.py
index f5024eb..92f37db 100644
--- a/pyanalyze/extensions.py
+++ b/pyanalyze/extensions.py
@@ -149,6 +149,27 @@ class NoAny(CustomCheck):
         )
 
 
+@dataclass(frozen=True)
+class ValidRegex(CustomCheck):
+    """Custom check that allows only values that are valid regular expressions.
+
+    Example::
+
+        def func(arg: Annotated[str, ValidRegex()]) -> None:
+            ...
+
+        func(".*")  # ok
+        func("[")  # error
+
+    """
+
+    def can_assign(self, value: "Value", ctx: "CanAssignContext") -> "CanAssign":
+        error = pyanalyze.implementation.check_regex_in_value(value)
+        if error is not None:
+            return error
+        return {}
+
+
 class _AsynqCallableMeta(type):
     def __getitem__(
         self, params: Tuple[Union[Literal[Ellipsis], List[object]], object]
diff --git a/pyanalyze/implementation.py b/pyanalyze/implementation.py
index 3676aa4..fd6cd4f 100644
--- a/pyanalyze/implementation.py
+++ b/pyanalyze/implementation.py
@@ -2,6 +2,7 @@ import ast
 import collections
 import collections.abc
 import inspect
+import re
 import typing
 from itertools import product
 from typing import (
@@ -20,6 +21,8 @@ from typing import (
 import qcore
 import typing_extensions
 
+import pyanalyze
+
 from . import runtime
 from .annotated_types import MaxLen, MinLen
 from .annotations import type_from_value
@@ -31,7 +34,9 @@ from .safe import hasattr_static, is_union, safe_isinstance, safe_issubclass
 from .signature import (
     ANY_SIGNATURE,
     CallContext,
+    ConcreteSignature,
     ImplReturn,
+    OverloadedSignature,
     ParameterKind,
     Signature,
     SigParameter,
@@ -2230,3 +2235,67 @@ def get_default_argspecs() -> Dict[object, Signature]:
             )
             signatures.append(sig)
     return {sig.callable: sig for sig in signatures}
+
+
+def check_regex(pattern: Union[str, bytes]) -> Optional[CanAssignError]:
+    try:
+        # TODO allow this without the useless isinstance()
+        if isinstance(pattern, str):
+            re.compile(pattern)
+        else:
+            re.compile(pattern)
+    except re.error as e:
+        return CanAssignError(
+            f"Invalid regex pattern: {e}", error_code=ErrorCode.invalid_regex
+        )
+    return None
+
+
+def check_regex_in_value(value: Value) -> Optional[CanAssignError]:
+    errors = []
+    for subval in flatten_values(value):
+        if not isinstance(subval, KnownValue):
+            continue
+        if not isinstance(subval.val, (str, bytes)):
+            continue
+        maybe_error = check_regex(subval.val)
+        if maybe_error is not None:
+            errors.append(maybe_error)
+    if errors:
+        if len(errors) == 1:
+            return errors[0]
+        return pyanalyze.value.CanAssignError(
+            "Invalid regex", errors, pyanalyze.error_code.ErrorCode.invalid_regex
+        )
+    return None
+
+
+def _re_impl_with_pattern(ctx: CallContext) -> Value:
+    pattern = ctx.vars["pattern"]
+    error = check_regex_in_value(pattern)
+    if error is not None:
+        ctx.show_error(error.message, error_code=ErrorCode.invalid_regex, arg="pattern")
+    return ctx.inferred_return_value
+
+
+def get_default_argspecs_with_cache(
+    asc: "pyanalyze.arg_spec.ArgSpecCache",
+) -> Dict[object, ConcreteSignature]:
+    sigs = {}
+    for func in (
+        re.compile,
+        re.search,
+        re.match,
+        re.fullmatch,
+        re.split,
+        re.findall,
+        re.finditer,
+        re.sub,
+        re.subn,
+    ):
+        sig = asc.get_argspec(func, impl=_re_impl_with_pattern)
+        assert isinstance(
+            sig, (Signature, OverloadedSignature)
+        ), f"failed to find signature for {func}: {sig}"
+        sigs[func] = sig
+    return sigs
diff --git a/pyanalyze/signature.py b/pyanalyze/signature.py
index d9f7da2..962675e 100644
--- a/pyanalyze/signature.py
+++ b/pyanalyze/signature.py
@@ -320,6 +320,8 @@ class CallContext:
     node: Optional[ast.AST]
     """AST node corresponding to the function call. Useful for
     showing errors."""
+    sig: "Signature"
+    inferred_return_value: Value
 
     def ast_for_arg(self, arg: str) -> Optional[ast.AST]:
         composite = self.composite_for_arg(arg)
@@ -1332,8 +1334,14 @@ class Signature:
                     visitor=ctx.visitor,
                     composites=composites,
                     node=ctx.node,
+                    sig=self,
+                    inferred_return_value=return_value,
                 )
-                return_value = self.impl(call_ctx)
+                with ctx.visitor.catch_errors() as caught_errors:
+                    return_value = self.impl(call_ctx)
+                if caught_errors:
+                    ctx.visitor.show_caught_errors(caught_errors)
+                    had_error = True
             elif self.evaluator is not None:
                 varmap = {
                     param: composite.value
@@ -1347,6 +1355,7 @@ class Signature:
                 )
                 return_value, errors = self.evaluator.evaluate(eval_ctx)
                 for error in errors:
+                    had_error = True
                     error_node = None
                     if error.argument is not None:
                         composite = bound_args[error.argument][1]


Test patch:
diff --git a/pyanalyze/test_implementation.py b/pyanalyze/test_implementation.py
index 930c96c..dea128e 100644
--- a/pyanalyze/test_implementation.py
+++ b/pyanalyze/test_implementation.py
@@ -1476,3 +1476,34 @@ class TestNamedTuple(TestNameCheckVisitorBase):
         def capybara() -> None:
             # on 3.13+ we get a second error from calling the runtime
             NamedTuple("x", None, y=int)  # E: incompatible_call  # E: incompatible_call
+
+
+class TestRegex(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_compile(self):
+        import re
+
+        def capybara():
+            re.compile("a")
+            re.compile(b"a")
+            re.compile("[")  # E: incompatible_call
+            re.compile(b"[")  # E: incompatible_call
+
+            re.sub(r"a", "b", "c")
+            re.sub(rb"(", b"b", b"c")  # E: incompatible_call
+
+            re.match(r"a", "b")
+            re.match(rb"(", b"b")  # E: incompatible_call
+
+    @assert_passes()
+    def test_extension(self):
+        from typing_extensions import Annotated
+
+        from pyanalyze.extensions import ValidRegex
+
+        def f(x: Annotated[str, ValidRegex()]):
+            pass
+
+        def capybara():
+            f("x")
+            f(r"[")  # E: invalid_regex