Files
2026-07-13 13:17:40 +08:00

89 lines
2.3 KiB
Python

from typing import Any
import pytest
from ray.dag.py_obj_scanner import _instances, _PyObjScanner
class Source:
pass
def test_simple_replace():
scanner = _PyObjScanner(source_type=Source)
my_objs = [Source(), [Source(), {"key": Source()}]]
found = scanner.find_nodes(my_objs)
assert len(found) == 3
replaced = scanner.replace_nodes({obj: 1 for obj in found})
assert replaced == [1, [1, {"key": 1}]]
def test_replace_multiple_types():
class OtherSource:
pass
scanner = _PyObjScanner(source_type=(Source, OtherSource))
my_objs = [Source(), [Source(), {"key": Source(), "key2": OtherSource()}]]
found = scanner.find_nodes(my_objs)
assert len(found) == 4
replaced = scanner.replace_nodes(
{obj: 1 if isinstance(obj, Source) else 2 for obj in found}
)
assert replaced == [1, [1, {"key": 1, "key2": 2}]]
def test_replace_nested_in_obj():
"""Test that the source can be nested in arbitrary objects."""
scanner = _PyObjScanner(source_type=Source)
class Outer:
def __init__(self, inner: Any):
self._inner = inner
def __eq__(self, other):
return self._inner == other._inner
my_objs = [Outer(Source()), Outer(Outer(Source())), Outer((Source(),))]
found = scanner.find_nodes(my_objs)
assert len(found) == 3
replaced = scanner.replace_nodes({obj: 1 for obj in found})
assert replaced == [Outer(1), Outer(Outer(1)), Outer((1,))]
def test_scanner_clear():
"""Test scanner clear to make the scanner GCable"""
prev_len = len(_instances)
def call_find_nodes():
scanner = _PyObjScanner(source_type=Source)
my_objs = [Source(), [Source(), {"key": Source()}]]
scanner.find_nodes(my_objs)
scanner.clear()
assert id(scanner) not in _instances
call_find_nodes()
assert prev_len == len(_instances)
def call_find_and_replace_nodes():
scanner = _PyObjScanner(source_type=Source)
my_objs = [Source(), [Source(), {"key": Source()}]]
found = scanner.find_nodes(my_objs)
scanner.replace_nodes({obj: 1 for obj in found})
scanner.clear()
assert id(scanner) not in _instances
call_find_and_replace_nodes()
assert prev_len == len(_instances)
if __name__ == "__main__":
import sys
sys.exit(pytest.main(["-v", __file__]))