File: test_tree.py

package info (click to toggle)
zarr 3.1.6-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 3,264 kB
  • sloc: python: 33,347; makefile: 11
file content (104 lines) | stat: -rw-r--r-- 3,181 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import textwrap
from typing import Any

import pytest

import zarr


@pytest.mark.parametrize("root_name", [None, "root"])
@pytest.mark.parametrize("atty", [True, False])
@pytest.mark.parametrize("plain", [True, False])
def test_tree(root_name: Any, atty: bool, plain: bool, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr("sys.stdout.isatty", lambda: atty)

    if atty and not plain:
        BOPEN = "\x1b[1m"
        BCLOSE = "\x1b[0m"
    else:
        BOPEN = ""
        BCLOSE = ""

    g = zarr.group(path=root_name)
    A = g.create_group("A")
    B = g.create_group("B")
    C = B.create_group("C")
    D = C.create_group("C")

    A.create_array(name="x", shape=(2), dtype="float64")
    A.create_array(name="y", shape=(0,), dtype="int8")
    B.create_array(name="x", shape=(0,), dtype="float64")
    C.create_array(name="x", shape=(0,), dtype="float64")
    D.create_array(name="x", shape=(0,), dtype="float64")

    result = repr(g.tree(plain=plain))
    root = root_name or ""

    expected = textwrap.dedent(f"""\
        {BOPEN}/{root}{BCLOSE}
        ├── {BOPEN}A{BCLOSE}
        │   ├── {BOPEN}x{BCLOSE} (2,) float64
        │   └── {BOPEN}y{BCLOSE} (0,) int8
        └── {BOPEN}B{BCLOSE}
            ├── {BOPEN}C{BCLOSE}
            │   ├── {BOPEN}C{BCLOSE}
            │   │   └── {BOPEN}x{BCLOSE} (0,) float64
            │   └── {BOPEN}x{BCLOSE} (0,) float64
            └── {BOPEN}x{BCLOSE} (0,) float64
        """)

    assert result == expected

    result = repr(g.tree(level=0, plain=plain))
    expected = textwrap.dedent(f"""\
        {BOPEN}/{root}{BCLOSE}
        ├── {BOPEN}A{BCLOSE}
        └── {BOPEN}B{BCLOSE}
        """)
    assert result == expected

    if not plain:
        tree = g.tree(plain=False)
        bundle = tree._repr_mimebundle_()
        assert "text/plain" in bundle
        assert "text/html" in bundle
        assert "<b>A</b>" in bundle["text/html"]
        assert "<b>x</b>" in bundle["text/html"]
        assert "<pre" in bundle["text/html"]


def test_tree_truncation() -> None:
    g = zarr.group()
    g.create_group("a")
    g.create_group("b")
    g.create_group("c")
    g.create_group("d")
    g.create_group("e")

    result = repr(g.tree(max_nodes=3, plain=True))
    assert "Truncated at max_nodes=3" in result
    # Should show exactly 3 nodes (lines with ── connectors).
    lines = result.strip().split("\n")
    node_lines = [line for line in lines if "──" in line]
    assert len(node_lines) == 3

    # Full tree should not show truncation message.
    full = repr(g.tree(max_nodes=500, plain=True))
    assert "truncated" not in full


def test_tree_html_escaping() -> None:
    g = zarr.group()
    g.create_group("<img onerror=alert(1) src=x>")

    tree = g.tree()
    bundle = tree._repr_mimebundle_()
    assert "&lt;img" in bundle["text/html"]
    assert "<img" not in bundle["text/html"]
    assert "<img onerror=alert(1) src=x>" in bundle["text/plain"]


def test_expand_not_implemented() -> None:
    g = zarr.group()
    with pytest.raises(NotImplementedError):
        g.tree(expand=True)