Skip to content

Commit

Permalink
Use an enum for sentinel
Browse files Browse the repository at this point in the history
  • Loading branch information
nicoddemus committed Sep 8, 2023
1 parent 3e95a37 commit 35771f6
Showing 1 changed file with 9 additions and 4 deletions.
13 changes: 9 additions & 4 deletions src/_pytest/assertion/rewrite.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
import tokenize
import types
from collections import defaultdict
from enum import Enum
from pathlib import Path
from pathlib import PurePath
from typing import Callable
Expand Down Expand Up @@ -53,7 +54,11 @@
PYC_EXT = ".py" + (__debug__ and "c" or "o")
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):
Expand Down Expand Up @@ -728,13 +733,13 @@ def run(self, mod: ast.Module) -> None:

# Collect asserts.
self.scope = (mod,)
nodes: List[Union[ast.AST, object]] = [mod]
nodes: List[Union[ast.AST, ScopeEndMarkerType]] = [mod]
while nodes:
node = nodes.pop()
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
self.scope = tuple((*self.scope, node))
nodes.append(_SCOPE_END_MARKER)
if node == _SCOPE_END_MARKER:
nodes.append(ScopeEndMarkerType.ScopeEndMarker)
if node is ScopeEndMarkerType.ScopeEndMarker:
self.scope = self.scope[:-1]
continue
assert isinstance(node, ast.AST)
Expand Down

0 comments on commit 35771f6

Please sign in to comment.