_rfile.py

# Author: Giacomo Parolini CERN 04/2025

r"""
\pythondoc RFile

TODO: document RFile

\code{.py}
# TODO code example
\endcode

\endpythondoc
"""

from . import pythonization


class _RFile_Get:
    """
    Allow access to objects through the method Get().
    This is pythonized to allow Get() to be called both with and without a template argument.
    """

    def __init__(self, rfile):
        self._rfile = rfile

    def __call__(self, namecycle):
        """
        Non-templated Get()
        """
        import ROOT

        key = self._rfile.GetKeyInfo(namecycle)
        if key:
            obj = ROOT.Internal.RFile_GetObjectFromKey(self._rfile, key)
            return ROOT._cppyy.bind_object(obj, key.GetClassName())
        # No key
        return None

    def __getitem__(self, template_arg):
        """
        Templated Get()
        """

        def getitem_wrapper(namecycle):
            obj = self._rfile._OriginalGet[template_arg](namecycle)
            return obj if obj else None

        return getitem_wrapper


class _RFile_Put:
    """
    Allow writing objects through the method Put().
    This is pythonized to allow Put() to be called both with and without a template argument.
    """

    def __init__(self, rfile):
        self._rfile = rfile

    def __call__(self, path, obj = None, title = ""):
        """
        Non-templated Put()
        """
        import ROOT

        if isinstance(path, ROOT.TNamed):
            # TNamed-specific overload: arguments are actually Put(namedObj, dirPath).
            path, obj = obj, path
            fullPath = ""
            if isinstance(path, str):
                fullPath = f"{path}/"
            elif path is not None:
                raise TypeError(f"Put(TNamed, dirPath) only accepts a string as the dirPath, not a {type(path)}")
            fullPath += obj.GetName()
            path = fullPath
            title = obj.GetTitle()

        objType = type(obj)
        if isinstance(obj, str):
            # special case: automatically convert python str to std::string
            className = "std::string"
        elif not hasattr(objType, "__cpp_name__"):
            raise TypeError(f"type {objType} is not supported by ROOT I/O")
        else:
            className = objType.__cpp_name__
        self._rfile.Put[className](path, obj, title)

    def __getitem__(self, template_arg):
        """
        Templated Put()
        """
        return self._rfile._OriginalPut[template_arg]


class _RFile_Overwrite:
    """
    Allow overwriting objects through the method Overwrite().
    This is pythonized to allow Overwrite() to be called both with and without a template argument.
    It also changes Overwrite's bitmask argument with keyword arguments.
    """

    def __init__(self, rfile):
        self._rfile = rfile

    def __call__(self, path, obj = None, **kwargs):
        """
        Non-templated Overwrite()
        """

        title = kwargs['title'] if 'title' in kwargs else ""
        createNewCycle = kwargs['createNewCycle'] if 'createNewCycle' in kwargs else True

        objType = type(obj)
        if isinstance(obj, str):
            # special case: automatically convert python str to std::string
            className = "std::string"
        elif not hasattr(objType, "__cpp_name__"):
            raise TypeError(f"type {objType} is not supported by ROOT I/O")
        else:
            className = objType.__cpp_name__

        from ROOT import RFile
        flags = (0 if createNewCycle else 1) * RFile.kOverwrite_NoNewCycle
        self._rfile.Overwrite[className](path, obj, title, flags)

    def __getitem__(self, template_arg):
        """
        Templated Overwrite()
        """
        return self._rfile._OriginalOverwrite[template_arg]


def _RFileExit(obj, exc_type, exc_val, exc_tb):
    """
    Close the RFile object.
    Signature and return value are imposed by Python, see
    https://docs.python.org/3/library/stdtypes.html#typecontextmanager.
    """
    obj.Close()
    return False


def _RFileOpen(original):
    """
    Pythonization for the factory methods (Recreate, Open, Update)
    """

    def rfile_open_wrapper(klass, *args):
        rfile = original(*args)
        rfile._OriginalGet = rfile.Get
        rfile.Get = _RFile_Get(rfile)
        rfile._OriginalPut = rfile.Put
        rfile.Put = _RFile_Put(rfile)
        rfile._OriginalOverwrite = rfile.Overwrite
        rfile.Overwrite = _RFile_Overwrite(rfile)
        return rfile

    return rfile_open_wrapper


def _RFileInit(rfile):
    """
    Prevent the creation of RFile through constructor (must use a factory method)
    """
    raise NotImplementedError("RFile can only be created via Recreate, Open or Update")


def _GetKeyInfo(rfile, path):
    key = rfile._OriginalGetKeyInfo(path)
    if key.has_value():
        return key.value()
    return None


def _ListKeys(rfile, basePath="", **kwargs):
    """
    Returns an iterable over all keys of objects and/or directories written into this RFile starting at path
    `basePath` (defaulting to include the content of all subdirectories).
    By default, keys referring to directories are not returned: only those referring to leaf objects are.
    If `basePath` is the path of a leaf object, only `basePath` itself will be returned.
    If it is the path of a directory, it will not be included in the listing.

    You can specify what to list via the keyword arguments:
    - if `listObjects == True`, the listing will include keys of non-directory objects (default);
    - if `listDirs == True`, the listing will include keys of directory objects;
    - if `listRecursive == True`, the listing will recurse on all subdirectories of `basePath` (default),
    otherwise it will only list immediate children of `basePath`.
        Example usage:
    ~~~{.py}
    for key in file.ListKeys():
        # iterate over all objects in the RFile
        print(f"{key.GetPath()};{key.GetCycle()} of type {key.GetClassName()}")

    for key in file.ListKeys("", listDirs=True):
        # iterate over all objects and directories in the RFile
        print(f"{key.GetPath()};{key.GetCycle()} of type {key.GetClassName()}")

    for key in file.ListKeys("a/b", listRecursive=False):
        # iterate over all objects that are immediate children of directory "a/b"
        print(f"{key.GetPath()};{key.GetCycle()} of type {key.GetClassName()}")

    for key in file.ListKeys("foo", listDirs=True, listObjects=False):
        # iterate over all directories under directory "foo", recursively
        print(key.GetPath())
    ~~~
    """
    from ROOT import RFile

    listObjects = kwargs['listObjects'] if 'listObjects' in kwargs else True
    listDirs = kwargs['listDirs'] if 'listDirs' in kwargs else False
    listRecursive = kwargs['listRecursive'] if 'listRecursive' in kwargs else True

    flags = (listObjects * RFile.kListObjects) | (listDirs * RFile.kListDirs) | (listRecursive * RFile.kListRecursive)
    iter = rfile._OriginalListKeys(basePath, flags)
    return iter


@pythonization("RFile", ns="ROOT")
def pythonize_rfile(klass):
    # Explicitly prevent to create a RFile via ctor
    klass.__init__ = _RFileInit

    # Pythonize factory methods
    klass.Open = classmethod(_RFileOpen(klass.Open))
    klass.Update = classmethod(_RFileOpen(klass.Update))
    klass.Recreate = classmethod(_RFileOpen(klass.Recreate))

    # Pythonization for __enter__ and __exit__ methods
    # These make RFile usable in a `with` statement as a context manager
    klass.__enter__ = lambda rfile: rfile
    klass.__exit__ = _RFileExit
    klass._OriginalGetKeyInfo = klass.GetKeyInfo
    klass.GetKeyInfo = _GetKeyInfo
    klass._OriginalListKeys = klass.ListKeys
    klass.ListKeys = _ListKeys