Use an enum for sentinel

This commit is contained in:
Bruno Oliveira 2023-09-08 18:44:12 -03:00
parent 3e95a370d0
commit 35771f643d
1 changed files with 9 additions and 4 deletions

View File

@ -14,6 +14,7 @@ import sys
import tokenize import tokenize
import types import types
from collections import defaultdict from collections import defaultdict
from enum import Enum
from pathlib import Path from pathlib import Path
from pathlib import PurePath from pathlib import PurePath
from typing import Callable from typing import Callable
@ -53,7 +54,11 @@ PYTEST_TAG = f"{sys.implementation.cache_tag}-pytest-{version}"
PYC_EXT = ".py" + (__debug__ and "c" or "o") PYC_EXT = ".py" + (__debug__ and "c" or "o")
PYC_TAIL = "." + PYTEST_TAG + PYC_EXT PYC_TAIL = "." + PYTEST_TAG + PYC_EXT
_SCOPE_END_MARKER = object()
class ScopeEndMarkerType(Enum):
"""Special marker that denotes we have just left a function or class definition."""
ScopeEndMarker = 1
class AssertionRewritingHook(importlib.abc.MetaPathFinder, importlib.abc.Loader): class AssertionRewritingHook(importlib.abc.MetaPathFinder, importlib.abc.Loader):
@ -728,13 +733,13 @@ class AssertionRewriter(ast.NodeVisitor):
# Collect asserts. # Collect asserts.
self.scope = (mod,) self.scope = (mod,)
nodes: List[Union[ast.AST, object]] = [mod] nodes: List[Union[ast.AST, ScopeEndMarkerType]] = [mod]
while nodes: while nodes:
node = nodes.pop() node = nodes.pop()
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
self.scope = tuple((*self.scope, node)) self.scope = tuple((*self.scope, node))
nodes.append(_SCOPE_END_MARKER) nodes.append(ScopeEndMarkerType.ScopeEndMarker)
if node == _SCOPE_END_MARKER: if node is ScopeEndMarkerType.ScopeEndMarker:
self.scope = self.scope[:-1] self.scope = self.scope[:-1]
continue continue
assert isinstance(node, ast.AST) assert isinstance(node, ast.AST)