_facade.py

import importlib
import os
import sys
import types
from functools import partial


class PyROOTConfiguration(object):
    """Class for configuring PyROOT"""

    def __init__(self):
        self.IgnoreCommandLineOptions = True
        self.ShutDown = True
        self.DisableRootLogon = False
        self.StartGUIThread = True


class _gROOTWrapper(object):
    """Internal class to manage lookups of gROOT in the facade.
    This wrapper calls _finalSetup on the facade when it
    receives a lookup, unless the lookup is for SetBatch.
    This allows to evaluate the command line parameters
    before checking if batch mode is on in _finalSetup
    """

    def __init__(self, facade):
        self.__dict__["_facade"] = facade

    @property
    def _gROOT(self):
        gROOT = self.__dict__.get("_gROOT")
        if gROOT is None:
            gROOT = self._facade._cppyy.gbl.ROOT.GetROOT()
            self.__dict__["_gROOT"] = gROOT
        return gROOT

    def __getattr__(self, name):
        if name != "SetBatch" and self._facade.__dict__["gROOT"] != self._gROOT:
            self._facade._finalSetup()
        return getattr(self._gROOT, name)

    def __setattr__(self, name, value):
        return setattr(self._gROOT, name, value)


class TDirectoryPythonAdapter:
    """Class analogous to the `ROOT::Internal::TDirectoryAtomicAdapter` on the
    C++ side, implementing the semantics that is expected from a
    `TDirectory *&` to the global directory (which is what gDirectory was
    originally), but in a thread-safe way.

    On the C++ side the following trick is used in TDirectory.h to achieve this
    We're re-implementing the expected semantics in Python, because the way how
    it is implemented in C++ is too contrived for it to get Pythonized
    automatically:

    ```C++
    #define gDirectory (::ROOT::Internal::TDirectoryAtomicAdapter{})
    ```

    So in C++, gDirectory is an adapter object that lazily converts to current
    TDirectory when you use it. It's implemented as as a preprocessor macro,
    which is then pretending to be a `TDirectory *` to the ROOT reflection system
    in TROOT.cxx:

    ```C++
    TGlobalMappedFunction::MakeFunctor("gDirectory", "TDirectory*", TDirectory::CurrentDirectory);
    ```

    This is quite hacky, and it is ambiguous to the the dynamic Python bindings
    layer at which point the implicit conversion to the current directory
    pointer should happen.

    For this reason, it's better to re-implement a specific adapter for Python,
    which skips all of that.

    Note: the C++ adapter also implements an assignment operator (operator=),
    but since this concept doesn't exist in Python outside data member
    re-assignment, it is not implemented here.
    """

    def _current_directory(self):
        import ROOT

        return ROOT.TDirectory.CurrentDirectory().load()

    def __getattr__(self, name):
        return getattr(self._current_directory(), name)

    def __str__(self):
        return str(self._current_directory())

    def __repr__(self):
        return repr(self._current_directory())

    def __eq__(self, other):
        import ROOT

        if other is self:
            return True
        cd = self._current_directory()
        if cd == ROOT.nullptr:
            return other == ROOT.nullptr
        return cd.IsEqual(other)

    def __ne__(self, other):
        import ROOT

        if other is self:
            return False
        cd = self._current_directory()
        if cd == ROOT.nullptr:
            return other != ROOT.nullptr
        return not cd.IsEqual(other)

    def __bool__(self):
        import ROOT

        return self._current_directory() != ROOT.nullptr

    def __cast_cpp__(self):
        """Casting to TDirectory for use in C++ functions."""
        return self._current_directory()


def _subimport(name):
    # type: (str) -> types.ModuleType
    """
    Import and return the Python module with the input name.

    Helper function for the __reduce__ method of the ROOTFacade class.
    """
    return importlib.import_module(name)


class ROOTFacade(types.ModuleType):
    """Facade class for ROOT module"""

    def __init__(self, module, is_ipython):
        types.ModuleType.__init__(self, module.__name__)

        self.__all__ = module.__all__
        self.__name__ = module.__name__
        self.__file__ = module.__file__
        self.__spec__ = module.__spec__
        self.__path__ = module.__path__
        self.__doc__ = module.__doc__
        self.__package__ = module.__package__
        self.__loader__ = module.__loader__

        # Inject gROOT global
        self.gROOT = _gROOTWrapper(self)

        # Inject the gDirectory adapter, mimicking the behavior of the
        # gDirectory preprocessor macro on the C++ side
        self.gDirectory = TDirectoryPythonAdapter()

        # Initialize configuration
        self.PyConfig = PyROOTConfiguration()

        self._is_ipython = is_ipython

        # Redirect lookups to temporary helper methods
        # This lets the user do some actions before all the machinery is in place:
        # - Set batch mode in gROOT
        # - Set options in PyConfig
        self.__class__.__getattr__ = self._getattr
        self.__class__.__setattr__ = self._setattr

    def SetHeuristicMemoryPolicy(self, enabled):
        import textwrap
        import warnings

        msg = """ROOT.SetHeuristicMemoryPolicy() is deprecated and will be removed in ROOT 6.44.
        Since ROOT 6.40, the heuristic memory policy is disabled by default, and with
        ROOT 6.44 it won't be possible to re-enable it with ROOT.SetHeuristicMemoryPolicy(),
        which was only meant to be used during a transition period and will be removed.
        """
        warnings.warn(textwrap.dedent(msg), FutureWarning, stacklevel=0)
        return self._cppyy._backend.SetHeuristicMemoryPolicy(enabled)

    def AddressOf(self, obj):
        # Return an indexable buffer of length 1, whose only element
        # is the address of the object.
        # The address of the buffer is the same as the address of the
        # address of the object

        # addr is the address of the address of the object
        addr = self.addressof(instance=obj, byref=True)

        # Check for 64 bit as suggested here:
        # https://docs.python.org/3/library/platform.html#cross-platform
        out_type = "Long64_t*" if sys.maxsize > 2**32 else "Int_t*"

        # Create a buffer (LowLevelView) from address
        return self._cppyy.ll.cast[out_type](addr)

    def _fallback_getattr(self, name):
        # Try:
        # - in the global namespace
        # - in the ROOT namespace
        # - in gROOT (ROOT lists such as list of files,
        #   memory mapped files, functions, geometries ecc.)
        # The first two attempts allow to lookup
        # e.g. ROOT.ROOT.Math as ROOT.Math

        # Note that hasattr caches the lookup for getattr
        if hasattr(self._cppyy.gbl, name):
            return getattr(self._cppyy.gbl, name)
        elif hasattr(self._cppyy.gbl.ROOT, name):
            return getattr(self._cppyy.gbl.ROOT, name)
        else:
            res = self.gROOT.FindObject(name)
            if res:
                return res
        raise AttributeError("Failed to get attribute {} from ROOT".format(name))

    def _register_converters_and_executors(self):
        converter_aliases = {
            "Float16_t": "float",
            "const Float16_t&": "const float&",
            "Double32_t": "double",
            "Double32_t&": "double&",
            "const Double32_t&": "const double&",
        }

        executor_aliases = {
            "Float16_t": "float",
            "Float16_t&": "float&",
            "Double32_t": "double",
            "Double32_t&": "double&",
        }

        from ROOT.libROOTPythonizations import CPyCppyyRegisterConverterAlias, CPyCppyyRegisterExecutorAlias

        for name, target in converter_aliases.items():
            CPyCppyyRegisterConverterAlias(name, target)

        for name, target in executor_aliases.items():
            CPyCppyyRegisterExecutorAlias(name, target)

    def _finalSetup(self):
        """
        Perform the final ROOT initialization.

        This method is intentionally deferred and is the *only* place where
        cppyy is imported and the C++ runtime is initialized. Delaying this
        step avoids importing the heavy-weight cppyy machinery unless it is
        actually required (for example, when accessing C++ ROOT symbols),
        allowing Python-only ROOT submodules to be imported with minimal
        overhead.
        """
        import cppyy
        import cppyy.ll
        import cppyy.types

        from ._application import PyROOTApplication
        from ._pythonization import _register_pythonizations, pythonization

        # signal policy: don't abort interpreter in interactive mode
        cppyy._backend.SetGlobalSignalPolicy(not cppyy.gbl.ROOT.GetROOT().IsBatch())

        self.__dict__["_cppyy"] = cppyy

        # Expose some functionality from CPyCppyy extension module
        cppyy_exports = [
            "nullptr",
            "bind_object",
            "as_cobject",
            "addressof",
            "SetImplicitSmartPointerConversion",
            "SetOwnership",
        ]
        for name in cppyy_exports:
            self.__dict__[name] = getattr(cppyy._backend, name)
        # For backwards compatibility
        self.__dict__["MakeNullPointer"] = partial(cppyy._backend.bind_object, 0)
        self.__dict__["BindObject"] = cppyy._backend.bind_object
        self.__dict__["AsCObject"] = cppyy._backend.as_cobject

        # Trigger the addition of the pythonizations
        _register_pythonizations()

        # Prevent this method from being re-entered through the gROOT wrapper
        self.__dict__["gROOT"] = self._cppyy.gbl.ROOT.GetROOT()

        # Make sure the interpreter is initialized once gROOT has been initialized
        self._cppyy.gbl.TInterpreter.Instance()

        # Setup interactive usage from Python
        self.__dict__["app"] = PyROOTApplication(self.PyConfig, self._is_ipython)
        if not self.gROOT.IsBatch() and self.PyConfig.StartGUIThread:
            self.app.init_graphics(self._cppyy.gbl.gEnv, self._cppyy.gbl.gSystem)

        # The automatic conversion of ordinary objects to smart pointers is
        # disabled for ROOT because it can cause trouble with overload
        # resolution. If a function has overloads for both ordinary objects and
        # smart pointers, then the implicit conversion to smart pointers can
        # result in the smart pointer overload being hit, even though there
        # would be an overload for the regular object. Since PyROOT didn't have
        # this feature before 6.32 anyway, disabling it was the safest option.
        self.SetImplicitSmartPointerConversion(False)

        # Redirect lookups to cppyy's global namespace
        self.__class__.__getattr__ = self._fallback_getattr
        self.__class__.__setattr__ = lambda self, name, val: setattr(self._cppyy.gbl, name, val)

        # Register custom converters and executors
        self._register_converters_and_executors()

        # Run rootlogon if exists
        self._run_rootlogon()

        # @pythonization decorator
        self.pythonization = pythonization

    def _getattr(self, name):
        # Special case, to allow "from ROOT import gROOT" w/o starting the graphics
        if name == "__path__":
            raise AttributeError(name)

        self._finalSetup()

        return getattr(self, name)

    def _setattr(self, name, val):
        # Setting attributes of ROOT will also define variables in the C++
        # runtime, so we generally require _finalSetup() in this method.
        #
        # Don't setup on submodule imports like `import ROOT._distrdf`,
        # implying a setattr(ROOT, "_distdf", <Module instance>) inside Pythons
        # importlib.
        if isinstance(val, types.ModuleType):
            self.__dict__[name] = val
            return

        self._finalSetup()

        return setattr(self, name, val)

    def _execute_rootlogon_module(self, file_path):
        """Execute the 'rootlogon.py' module found at the given 'file_path'"""
        # Could also have used execfile, but import is likely to give fewer surprises
        module_name = "rootlogon"

        import importlib.util

        spec = importlib.util.spec_from_file_location(module_name, file_path)
        module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = module
        spec.loader.exec_module(module)

    def _run_rootlogon(self):
        # Run custom logon file (must be after creation of ROOT globals)
        hasargv = hasattr(sys, "argv")
        # -n disables the reading of the logon file, just like with root
        if hasargv and "-n" not in sys.argv and not self.PyConfig.DisableRootLogon:
            file_path_home = os.path.expanduser("~/.rootlogon.py")
            file_path_local = os.path.join(os.getcwd(), ".rootlogon.py")
            if os.path.exists(file_path_home):
                self._execute_rootlogon_module(file_path_home)
            elif os.path.exists(file_path_local):
                self._execute_rootlogon_module(file_path_local)
            else:
                # If the .py version of rootlogon exists, the .C is ignored (the user can
                # load the .C from the .py, if so desired).
                # System logon, user logon, and local logon (skip Rint.Logon)
                name = ".rootlogon.C"
                logons = [
                    os.path.join(str(self.TROOT.GetEtcDir()), "system" + name),
                    os.path.expanduser(os.path.join("~", name)),
                ]
                if logons[-1] != os.path.join(os.getcwd(), name):
                    logons.append(name)
                for rootlogon in logons:
                    if os.path.exists(rootlogon):
                        self.TApplication.ExecuteFile(rootlogon)

    def __reduce__(self):
        # type: () -> types.ModuleType
        """
        Reduction function of the ROOT facade to customize the (pickle)
        serialization step.

        Defines the ingredients needed for a correct serialization of the
        facade, that is a function that imports a Python module and the name of
        that module, which corresponds to this facade's __name__ attribute. This
        method helps serialization tools like `cloudpickle`, especially used in
        distributed environments, that always need to include information about
        the ROOT module in the serialization step. For example, the following
        snippet would not work without this method::

            import ROOT
            import cloudpickle

            def foo():
                return ROOT.TH1F()

            cloudpickle.loads(cloudpickle.dumps(foo))

        In particular, it would raise::

            TypeError: cannot pickle 'ROOTFacade' object
        """
        return _subimport, (self.__name__,)

    # Inject version as __version__ property in ROOT module
    @property
    def __version__(self):
        return self.gROOT.GetVersion()

    # Overload VecOps namespace
    # The property gets the C++ namespace, adds the pythonizations and
    # eventually deletes itself so that following calls go directly
    # to the C++ namespace. This mechanic ensures that we pythonize the
    # namespace lazily.
    @property
    def VecOps(self):
        ns = self._fallback_getattr("VecOps")
        from ._pythonization._rvec import _AsRVec

        ns.AsRVec = _AsRVec
        del type(self).VecOps
        return ns

    # Overload RDF namespace
    @property
    def RDF(self):
        self._finalSetup()
        ns = self._fallback_getattr("RDF")

        def MakeCSVDataFrame(fileName, readHeaders=True, delimiter=",", linesChunkSize=-1, colTypes={}, **kwargs):
            options = ns.RCsvDS.ROptions()
            options.fHeaders = readHeaders
            options.fDelimiter = delimiter
            options.fLinesChunkSize = linesChunkSize
            options.fColumnTypes = colTypes
            for key, val in kwargs.items():
                structMemberName = "f" + key[0].upper() + key[1:]
                if hasattr(options, structMemberName):
                    setattr(options, structMemberName, val)
            return ns._FromCSV(fileName, options)

        if hasattr(ns, "FromCSV"):
            # Provide a FromCSV factory method that uses keyword arguments instead of the ROptions config struct.
            # In Python, the RCsvDS::ROptions struct members are available without the leading 'f' and in camelCase,
            # e.g. fDelimiter --> delimiter.
            # We need to keep the parameters of the old FromCSV signature for backward compatibility.
            ns._FromCSV = ns.FromCSV
            ns.FromCSV = MakeCSVDataFrame

        # Make a copy of the arrays that have strides to make sure we read the correct values
        # TODO a cleaner fix
        def MakeNumpyDataFrameCopy(np_dict):
            import numpy

            from ._pythonization._rdataframe import _MakeNumpyDataFrame

            for key in np_dict.keys():
                if (np_dict[key].__array_interface__["strides"]) is not None:
                    np_dict[key] = numpy.copy(np_dict[key])
            return _MakeNumpyDataFrame(np_dict)

        ns.FromNumpy = MakeNumpyDataFrameCopy

        # make a RDataFrame from a Pandas dataframe
        def MakePandasDataFrame(df):
            from ._pythonization._rdataframe import _MakeNumpyDataFrame

            np_dict = {}
            for key in df.columns:
                np_dict[key] = df[key].to_numpy()
            return _MakeNumpyDataFrame(np_dict)

        ns.FromPandas = MakePandasDataFrame

        try:
            # Inject Pythonizations to interact between local and distributed RDF package
            from ._pythonization._rdf_namespace import (
                _create_distributed_module,
                _fromspec,
                _rungraphs,
                _variationsfor,
            )

            ns.Distributed = _create_distributed_module(ns)
            # Inject the experimental package which shows a warning before usage
            ns.Experimental.Distributed = _create_distributed_module(ns, True)
            ns.RunGraphs = _rungraphs(ns.Distributed.RunGraphs, ns.RunGraphs)
            ns.Experimental.VariationsFor = _variationsfor(ns.Distributed.VariationsFor, ns.Experimental.VariationsFor)
            ns.Experimental.FromSpec = _fromspec(ns.Distributed.FromSpec, ns.Experimental.FromSpec)
        except ImportError:
            # _rdf_namespace submodule not available (expected for dataframe=OFF)
            pass

        del type(self).RDF
        return ns

    @property
    def RDataFrame(self):
        """
        Dispatch between the local and distributed RDataFrame depending on
        input arguments.
        """
        local_rdf = self.__getattr__("RDataFrame")
        try:
            import ROOT._distrdf

            from ._pythonization._rdf_namespace import _rdataframe

            return _rdataframe(local_rdf, ROOT._distrdf.RDataFrame)
        except ImportError:
            # _distrdf submodule not available (expected for dataframe=OFF)
            return local_rdf

    # Overload RooFit namespace
    @property
    def RooFit(self):
        ns = self._fallback_getattr("RooFit")
        try:
            from ._pythonization._roofit import pythonize_roofit_namespace
        except ImportError:
            # _roofit submodule not available (expected for roofit=OFF)
            del type(self).RooFit
            return ns

        pythonize_roofit_namespace(ns)

        del type(self).RooFit
        return ns

    # Overload TMVA namespace
    @property
    def TMVA(self):
        ns = self._fallback_getattr("TMVA")
        try:
            # This line is needed to import the pythonizations in _tmva directory.
            # The comment suppresses linter errors about unused imports.
            from ._pythonization import _tmva  # noqa: F401
            from ._pythonization._tmva._rtensor import _AsRTensor
            from ._pythonization._tmva._sofie._parser._keras.parser import PyKeras
            from ._pythonization._tmva._tree_inference import SaveXGBoost

            setattr(ns.Experimental.SOFIE, "PyKeras", PyKeras)

            ns.Experimental.AsRTensor = _AsRTensor
            ns.Experimental.SaveXGBoost = SaveXGBoost
        except ImportError:
            # _tmva submodule not available (expected for tmva=OFF)
            pass
        del type(self).TMVA
        return ns

    # Create and overload Numba namespace
    @property
    def Numba(self):
        from ._numbadeclare import _NumbaDeclareDecorator

        self._cppyy.cppdef("namespace Numba {}")
        ns = self._fallback_getattr("Numba")
        ns.Declare = staticmethod(_NumbaDeclareDecorator)
        del type(self).Numba
        return ns

    @property
    def NumbaExt(self):
        import numba

        if not hasattr(numba, "version_info") or numba.version_info < (0, 54):
            raise Exception("NumbaExt requires Numba version 0.54 or higher")

        # The comment in the next line suppresses linter errors about unused imports
        import cppyy.numba_ext  # noqa: F401

        # Return something as it is a property function
        return self

    # Get TPyDispatcher for programming GUI callbacks
    @property
    def TPyDispatcher(self):
        self._cppyy.include("ROOT/TPyDispatcher.h")
        tpd = self._cppyy.gbl.TPyDispatcher
        type(self).TPyDispatcher = tpd
        return tpd

    # Create the uhi namespace
    @property
    def uhi(self):
        uhi_module = types.ModuleType("uhi")
        uhi_module.__file__ = "<module ROOT>"
        uhi_module.__package__ = self
        from ._pythonization._uhi import _add_module_level_uhi_helpers

        _add_module_level_uhi_helpers(uhi_module)
        return uhi_module

    @property
    def Experimental(self):
        ns = self._fallback_getattr("Experimental")

        from ._pythonization._ml_dataloader import _inject_dataloader_api

        _inject_dataloader_api(ns.ML)

        return ns