_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)


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.__cached__ = module.__cached__
        self.__path__ = module.__path__
        self.__doc__ = module.__doc__
        self.__package__ = module.__package__
        self.__loader__ = module.__loader__

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

        # 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 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

        self.__dict__["_cppyy"] = cppyy

        # Expose some functionality from CPyCppyy extension module
        cppyy_exports = [
            "nullptr",
            "bind_object",
            "as_cobject",
            "addressof",
            "SetHeuristicMemoryPolicy",
            "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)

        # Set memory policy to kUseHeuristics.
        # This restores the default in PyROOT which was changed
        # by new Cppyy
        self.SetHeuristicMemoryPolicy(True)

        # The automatic conversion of ordinary obejcts 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):
        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")
        try:
            from ._pythonization._rvec import _AsRVec

            ns.AsRVec = _AsRVec
        except Exception:
            raise Exception("Failed to pythonize the namespace VecOps")
        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:
            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 DistRDF

            from ._pythonization._rdf_namespace import _rdataframe

            return _rdataframe(local_rdf, DistRDF.RDataFrame)
        except ImportError:
            return local_rdf

    # Overload RooFit namespace
    @property
    def RooFit(self):
        from ._pythonization._roofit import pythonize_roofit_namespace

        ns = self._fallback_getattr("RooFit")
        try:
            pythonize_roofit_namespace(ns)
        except Exception:
            raise Exception("Failed to pythonize the namespace RooFit")
        del type(self).RooFit
        return ns

    # Overload TMVA namespace
    @property
    def TMVA(self):
        # 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

        ns = self._fallback_getattr("TMVA")
        hasRDF = "dataframe" in self.gROOT.GetConfigFeatures()
        if hasRDF:
            try:
                from ._pythonization._tmva import inject_rbatchgenerator
                from ._pythonization._tmva._rtensor import _AsRTensor
                from ._pythonization._tmva._tree_inference import SaveXGBoost

                inject_rbatchgenerator(ns)
                ns.Experimental.AsRTensor = _AsRTensor
                ns.Experimental.SaveXGBoost = SaveXGBoost
            except Exception:
                raise Exception("Failed to pythonize the namespace TMVA")
        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
        try:
            from ._pythonization._uhi import _add_module_level_uhi_helpers

            _add_module_level_uhi_helpers(uhi_module)
        except ImportError:
            raise Exception("Failed to pythonize the namespace uhi")
        return uhi_module