expression.py 9.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331
  1. r"""Evaluate match expressions, as used by `-k` and `-m`.
  2. The grammar is:
  3. expression: expr? EOF
  4. expr: and_expr ('or' and_expr)*
  5. and_expr: not_expr ('and' not_expr)*
  6. not_expr: 'not' not_expr | '(' expr ')' | ident kwargs?
  7. ident: (\w|:|\+|-|\.|\[|\]|\\|/)+
  8. kwargs: ('(' name '=' value ( ', ' name '=' value )* ')')
  9. name: a valid ident, but not a reserved keyword
  10. value: (unescaped) string literal | (-)?[0-9]+ | 'False' | 'True' | 'None'
  11. The semantics are:
  12. - Empty expression evaluates to False.
  13. - ident evaluates to True or False according to a provided matcher function.
  14. - or/and/not evaluate according to the usual boolean semantics.
  15. - ident with parentheses and keyword arguments evaluates to True or False according to a provided matcher function.
  16. """
  17. from __future__ import annotations
  18. import ast
  19. from collections.abc import Iterator
  20. from collections.abc import Mapping
  21. from collections.abc import Sequence
  22. import dataclasses
  23. import enum
  24. import keyword
  25. import re
  26. import types
  27. from typing import Literal
  28. from typing import NoReturn
  29. from typing import overload
  30. from typing import Protocol
  31. __all__ = [
  32. "Expression",
  33. "ParseError",
  34. ]
  35. class TokenType(enum.Enum):
  36. LPAREN = "left parenthesis"
  37. RPAREN = "right parenthesis"
  38. OR = "or"
  39. AND = "and"
  40. NOT = "not"
  41. IDENT = "identifier"
  42. EOF = "end of input"
  43. EQUAL = "="
  44. STRING = "string literal"
  45. COMMA = ","
  46. @dataclasses.dataclass(frozen=True)
  47. class Token:
  48. __slots__ = ("pos", "type", "value")
  49. type: TokenType
  50. value: str
  51. pos: int
  52. class ParseError(Exception):
  53. """The expression contains invalid syntax.
  54. :param column: The column in the line where the error occurred (1-based).
  55. :param message: A description of the error.
  56. """
  57. def __init__(self, column: int, message: str) -> None:
  58. self.column = column
  59. self.message = message
  60. def __str__(self) -> str:
  61. return f"at column {self.column}: {self.message}"
  62. class Scanner:
  63. __slots__ = ("current", "tokens")
  64. def __init__(self, input: str) -> None:
  65. self.tokens = self.lex(input)
  66. self.current = next(self.tokens)
  67. def lex(self, input: str) -> Iterator[Token]:
  68. pos = 0
  69. while pos < len(input):
  70. if input[pos] in (" ", "\t"):
  71. pos += 1
  72. elif input[pos] == "(":
  73. yield Token(TokenType.LPAREN, "(", pos)
  74. pos += 1
  75. elif input[pos] == ")":
  76. yield Token(TokenType.RPAREN, ")", pos)
  77. pos += 1
  78. elif input[pos] == "=":
  79. yield Token(TokenType.EQUAL, "=", pos)
  80. pos += 1
  81. elif input[pos] == ",":
  82. yield Token(TokenType.COMMA, ",", pos)
  83. pos += 1
  84. elif (quote_char := input[pos]) in ("'", '"'):
  85. end_quote_pos = input.find(quote_char, pos + 1)
  86. if end_quote_pos == -1:
  87. raise ParseError(
  88. pos + 1,
  89. f'closing quote "{quote_char}" is missing',
  90. )
  91. value = input[pos : end_quote_pos + 1]
  92. if (backslash_pos := input.find("\\")) != -1:
  93. raise ParseError(
  94. backslash_pos + 1,
  95. r'escaping with "\" not supported in marker expression',
  96. )
  97. yield Token(TokenType.STRING, value, pos)
  98. pos += len(value)
  99. else:
  100. match = re.match(r"(:?\w|:|\+|-|\.|\[|\]|\\|/)+", input[pos:])
  101. if match:
  102. value = match.group(0)
  103. if value == "or":
  104. yield Token(TokenType.OR, value, pos)
  105. elif value == "and":
  106. yield Token(TokenType.AND, value, pos)
  107. elif value == "not":
  108. yield Token(TokenType.NOT, value, pos)
  109. else:
  110. yield Token(TokenType.IDENT, value, pos)
  111. pos += len(value)
  112. else:
  113. raise ParseError(
  114. pos + 1,
  115. f'unexpected character "{input[pos]}"',
  116. )
  117. yield Token(TokenType.EOF, "", pos)
  118. @overload
  119. def accept(self, type: TokenType, *, reject: Literal[True]) -> Token: ...
  120. @overload
  121. def accept(
  122. self, type: TokenType, *, reject: Literal[False] = False
  123. ) -> Token | None: ...
  124. def accept(self, type: TokenType, *, reject: bool = False) -> Token | None:
  125. if self.current.type is type:
  126. token = self.current
  127. if token.type is not TokenType.EOF:
  128. self.current = next(self.tokens)
  129. return token
  130. if reject:
  131. self.reject((type,))
  132. return None
  133. def reject(self, expected: Sequence[TokenType]) -> NoReturn:
  134. raise ParseError(
  135. self.current.pos + 1,
  136. "expected {}; got {}".format(
  137. " OR ".join(type.value for type in expected),
  138. self.current.type.value,
  139. ),
  140. )
  141. # True, False and None are legal match expression identifiers,
  142. # but illegal as Python identifiers. To fix this, this prefix
  143. # is added to identifiers in the conversion to Python AST.
  144. IDENT_PREFIX = "$"
  145. def expression(s: Scanner) -> ast.Expression:
  146. if s.accept(TokenType.EOF):
  147. ret: ast.expr = ast.Constant(False)
  148. else:
  149. ret = expr(s)
  150. s.accept(TokenType.EOF, reject=True)
  151. return ast.fix_missing_locations(ast.Expression(ret))
  152. def expr(s: Scanner) -> ast.expr:
  153. ret = and_expr(s)
  154. while s.accept(TokenType.OR):
  155. rhs = and_expr(s)
  156. ret = ast.BoolOp(ast.Or(), [ret, rhs])
  157. return ret
  158. def and_expr(s: Scanner) -> ast.expr:
  159. ret = not_expr(s)
  160. while s.accept(TokenType.AND):
  161. rhs = not_expr(s)
  162. ret = ast.BoolOp(ast.And(), [ret, rhs])
  163. return ret
  164. def not_expr(s: Scanner) -> ast.expr:
  165. if s.accept(TokenType.NOT):
  166. return ast.UnaryOp(ast.Not(), not_expr(s))
  167. if s.accept(TokenType.LPAREN):
  168. ret = expr(s)
  169. s.accept(TokenType.RPAREN, reject=True)
  170. return ret
  171. ident = s.accept(TokenType.IDENT)
  172. if ident:
  173. name = ast.Name(IDENT_PREFIX + ident.value, ast.Load())
  174. if s.accept(TokenType.LPAREN):
  175. ret = ast.Call(func=name, args=[], keywords=all_kwargs(s))
  176. s.accept(TokenType.RPAREN, reject=True)
  177. else:
  178. ret = name
  179. return ret
  180. s.reject((TokenType.NOT, TokenType.LPAREN, TokenType.IDENT))
  181. BUILTIN_MATCHERS = {"True": True, "False": False, "None": None}
  182. def single_kwarg(s: Scanner) -> ast.keyword:
  183. keyword_name = s.accept(TokenType.IDENT, reject=True)
  184. if not keyword_name.value.isidentifier():
  185. raise ParseError(
  186. keyword_name.pos + 1,
  187. f"not a valid python identifier {keyword_name.value}",
  188. )
  189. if keyword.iskeyword(keyword_name.value):
  190. raise ParseError(
  191. keyword_name.pos + 1,
  192. f"unexpected reserved python keyword `{keyword_name.value}`",
  193. )
  194. s.accept(TokenType.EQUAL, reject=True)
  195. if value_token := s.accept(TokenType.STRING):
  196. value: str | int | bool | None = value_token.value[1:-1] # strip quotes
  197. else:
  198. value_token = s.accept(TokenType.IDENT, reject=True)
  199. if (number := value_token.value).isdigit() or (
  200. number.startswith("-") and number[1:].isdigit()
  201. ):
  202. value = int(number)
  203. elif value_token.value in BUILTIN_MATCHERS:
  204. value = BUILTIN_MATCHERS[value_token.value]
  205. else:
  206. raise ParseError(
  207. value_token.pos + 1,
  208. f'unexpected character/s "{value_token.value}"',
  209. )
  210. ret = ast.keyword(keyword_name.value, ast.Constant(value))
  211. return ret
  212. def all_kwargs(s: Scanner) -> list[ast.keyword]:
  213. ret = [single_kwarg(s)]
  214. while s.accept(TokenType.COMMA):
  215. ret.append(single_kwarg(s))
  216. return ret
  217. class MatcherCall(Protocol):
  218. def __call__(self, name: str, /, **kwargs: str | int | bool | None) -> bool: ...
  219. @dataclasses.dataclass
  220. class MatcherNameAdapter:
  221. matcher: MatcherCall
  222. name: str
  223. def __bool__(self) -> bool:
  224. return self.matcher(self.name)
  225. def __call__(self, **kwargs: str | int | bool | None) -> bool:
  226. return self.matcher(self.name, **kwargs)
  227. class MatcherAdapter(Mapping[str, MatcherNameAdapter]):
  228. """Adapts a matcher function to a locals mapping as required by eval()."""
  229. def __init__(self, matcher: MatcherCall) -> None:
  230. self.matcher = matcher
  231. def __getitem__(self, key: str) -> MatcherNameAdapter:
  232. return MatcherNameAdapter(matcher=self.matcher, name=key[len(IDENT_PREFIX) :])
  233. def __iter__(self) -> Iterator[str]:
  234. raise NotImplementedError()
  235. def __len__(self) -> int:
  236. raise NotImplementedError()
  237. class Expression:
  238. """A compiled match expression as used by -k and -m.
  239. The expression can be evaluated against different matchers.
  240. """
  241. __slots__ = ("code",)
  242. def __init__(self, code: types.CodeType) -> None:
  243. self.code = code
  244. @classmethod
  245. def compile(cls, input: str) -> Expression:
  246. """Compile a match expression.
  247. :param input: The input expression - one line.
  248. """
  249. astexpr = expression(Scanner(input))
  250. code: types.CodeType = compile(
  251. astexpr,
  252. filename="<pytest match expression>",
  253. mode="eval",
  254. )
  255. return Expression(code)
  256. def evaluate(self, matcher: MatcherCall) -> bool:
  257. """Evaluate the match expression.
  258. :param matcher:
  259. Given an identifier, should return whether it matches or not.
  260. Should be prepared to handle arbitrary strings as input.
  261. :returns: Whether the expression matches or not.
  262. """
  263. ret: bool = bool(eval(self.code, {"__builtins__": {}}, MatcherAdapter(matcher)))
  264. return ret