Modify class instances inside a context mananger and keep track of them

I'm designing a small workflow manager similar to Apache Airflow where I want to build a directed computational graph. Every Node will inherit from Node and the graph is build inside a contextmanager dag. I'd like to

1. keep track of all instances of Node that are created inside the dag.
2. want the Nodes to know if there inside with dag() or not.

There would be two approaches to this:

1. I create a DAG instance and pass it to my Nodes upon creation. This would be very simple but the __init__ would have to accept the dag as kwargs.

dag = DAG()
Node(dag, **kwargs)

2. In the style of Apache Airflow using a context manager:

with dag():
    Node()

where I define a global list of all Nodes dags = [] as wel开发者_开发百科l as the parent Node and the context manager as follows:

import contextlib

class Node:
    _dag_ = None
    def __init__(self):
        if self._dag_ is not None:
            self._dag_.append(self)
            print("Inside DAG manager")
        else:
            print("Not inside DAG manager")  


dags = []

@contextlib.contextmanager
def dag():
    Node._dag_ = dags
    try:
        yield
    finally:
        Node._dag_ = None    

This works but looks a little bit scary to me. There would be issues with multithreading but I have no plans in supporting that. Is there a better solution to what I'm trying to do? Could there be any issues I'm missing? I tried to read through the Airflow code but couldn't understand how they do it. Does anyone know / understand?