pytorch-lightning

1
# Copyright The Lightning AI team.
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14
"""Utilities for traversing the tree of components in an app."""
15

16
from typing import TYPE_CHECKING, Type
17

18
import lightning.app
19

20
if TYPE_CHECKING:
21
    from lightning.app.utilities.types import Component, ComponentTuple
22

23

24
def breadth_first(root: "Component", types: Type["ComponentTuple"] = None):
25
    """Returns a generator that walks through the tree of components breadth-first.
26

27
    Arguments:
28
        root: The root component of the tree
29
        types: If provided, only the component types in this list will be visited.
30

31
    """
32
    yield from _BreadthFirstVisitor(root, types)
33

34

35
class _BreadthFirstVisitor:
36
    def __init__(self, root: "Component", types: Type["ComponentTuple"] = None) -> None:
37
        self.queue = [root]
38
        self.types = types
39

40
    def __iter__(self):
41
        return self
42

43
    def __next__(self):
44
        from lightning.app.structures import Dict
45

46
        while self.queue:
47
            component = self.queue.pop(0)
48

49
            if isinstance(component, lightning.app.LightningFlow):
50
                components = [getattr(component, el) for el in sorted(component._flows)]
51
                for struct_name in sorted(component._structures):
52
                    structure = getattr(component, struct_name)
53
                    if isinstance(structure, Dict):
54
                        values = sorted(structure.items(), key=lambda x: x[0])
55
                    else:
56
                        values = sorted(((v.name, v) for v in structure), key=lambda x: x[0])
57
                    for _, value in values:
58
                        if isinstance(value, lightning.app.LightningFlow):
59
                            components.append(value)
60
                self.queue += components
61
                self.queue += component.works(recurse=False)
62

63
            if any(isinstance(component, t) for t in self.types):
64
                return component
65

66
        raise StopIteration
67

68

69
class _DepthFirstVisitor:
70
    def __init__(self, root: "Component", types: Type["ComponentTuple"] = None) -> None:
71
        self.stack = [root]
72
        self.types = types
73

74
    def __iter__(self):
75
        return self
76

77
    def __next__(self):
78
        while self.stack:
79
            component = self.stack.pop()
80

81
            if isinstance(component, lightning.app.LightningFlow):
82
                self.stack += list(component.flows.values())
83
                self.stack += component.works(recurse=False)
84

85
            if any(isinstance(component, t) for t in self.types):
86
                return component
87

88
        raise StopIteration
89