Source code for est.tests.test_prettry_print

from typing import List
from typing import Optional

import numpy
from pydantic import BaseModel

from ..core.utils import pretty_print




class InnerModel(BaseModel, arbitrary_types_allowed=True):
    numbers: numpy.ndarray
    optional_value: Optional[int]





class OuterModel(BaseModel):
    name: str
    inner: InnerModel
    attributes: dict
    tags: List[str]





def test_tree_complex_case():
    data = OuterModel(
        name="TestObject",
        inner=InnerModel(
            numbers=numpy.array([[1, 2, 3], [4, 5, 6]]),
            optional_value=None,
        ),
        attributes={"key1": "value1", "key2": 42},
        tags=["a", "b", "c"],
    )

    result = pretty_print.tree(data, title="RootObject")

    expected_unicode = """RootObject
├── name: 'TestObject'
├── inner
│   ├── numbers: shape=(2, 3) min=1 max=6
│   └── optional_value: <MISSING>
├── attributes
│   ├── key1: 'value1'
│   └── key2: 42
└── tags: len=3"""

    expected_ascii = """RootObject
|-- name: 'TestObject'
|-- inner
|   |-- numbers: shape=(2, 3) min=1 max=6
|   `-- optional_value: <MISSING>
|-- attributes
|   |-- key1: 'value1'
|   `-- key2: 42
`-- tags: len=3"""

    expected = expected_unicode if pretty_print._supports_unicode() else expected_ascii

    assert result == expected