## \file
## \ingroup tutorial_io
## \notebook -nodraw
## This tutorial demonstrates the usage of the TFile class as a Python context
## manager.
##
## \macro_code
## \macro_output
##
## \date March 2022
## \author Vincenzo Eduardo Padulano CERN/UPV
import os

import ROOT
from ROOT import TFile, gROOT

# By default, in ROOT 6, objects of some ROOT types such as `TH1` and its derived types
# are automatically attached to a ROOT.TDirectory when they are created.
# Objects are attached to the current directory denoted by ROOT.gDirectory.
# For ROOT 7, the default will be *not* to attach objects to ROOT.gDirectory, except for
# cases where the documentation states it explicitly, such as TTree.
# In ROOT 6, this mode can be tested by setting ROOT.Experimental.DisableObjectAutoRegistration().

# The next line will print 'PyROOT' as the name of the current directory.
# That is the global directory created when using ROOT from Python, which is
# the ROOT.gROOT object.
print("Current directory: '{}'.\n".format(ROOT.gDirectory.GetName()))

# We can check to which directory a newly created histogram is attached.
histo_1 = ROOT.TH1F("histo_1", "histo_1", 10, 0, 10)
print(
    "Histogram '{}' is attached to: '{}'.\n".format(
        histo_1.GetName(), "Nothing" if not histo_1.GetDirectory() else histo_1.GetDirectory().GetName()
    )
)

# For quick saving and forgetting of objects into ROOT files, it is possible to
# open a TFile as a Python context manager. In the context, objects can be
# created, modified and finally written to the file. At the end of the context,
# the file will be automatically closed.
path = str(gROOT.GetTutorialDir()) + '/io/'
filename = path+"tfile_1.root"
with TFile.Open(filename, "recreate") as f:
    histo_2 = ROOT.TH1F("histo_2", "histo_2", 10, 0, 10)
    # Inside the context, the current directory is the open file
    print("Current directory: '{}'.\n".format(ROOT.gDirectory.GetName()))
    # In ROOT 6, the created histogram is automatically attached to the file
    print(
        "Histogram '{}' is attached to: '{}'.\n".format(
            histo_2.GetName(), histo_2.GetDirectory().GetName() if histo_2.GetDirectory() else "Nothing"
        )
    )
    # When auto registration is off, this can be done explicitly:
    histo_2.SetDirectory(f)
    print(
        "After setting the directory, histogram '{}' is attached to: '{}'.\n".format(
            histo_2.GetName(), histo_2.GetDirectory().GetName()
        )
    )
    # Before exiting the context, objects can be written to the file
    f.WriteObject(histo_2, "my_histogram")

# When the TFile.Close method is called, the current directory is automatically
# set again to ROOT.gROOT. Objects that were attached to the file inside the
# context are automatically deleted, so they can't be accessed anymore after
# the file is closed.
print("Status after the first TFile context manager:")
print(" Current directory: '{}'.".format(ROOT.gDirectory.GetName()))
# print(histo_2) # the object is deleted at this point, so don't use it anymore!

# Also reading data from a TFile can be done in a context manager. Information
# stored in the objects of the file can be queried and used inside the context.
# After the context, the objects are not usable anymore because the file is
# automatically closed. This means you should use this pattern as a quick way
# to get information or modify objects from a certain file, without needing to
# keep the histograms alive afterwards.

with TFile.Open(filename, "read") as f:
    # Retrieve histogram using the name given to f.WriteObject in the previous
    # with statement
    histo_2_fromfile = f["my_histogram"]
    print("Retrieved '{}' histogram from file '{}'.\n".format(histo_2_fromfile.GetName(), f.GetName()))

# Cleanup the file created for this tutorial
os.remove(filename)
