Find runs on GRID directly with the Dirac python API (#90)

* Add method to fetch data from GRID with DIRAC API
(to avoid reading the NectarCAM ELog). The method using the Elog is kept as fallback.

* fix a strange bug where DIRAC modify environment
of logging which becomes in conflict with ctapipe

---------

Co-authored-by: guillaume.grolleron <[email protected]>

src/nectarchain/data/management.py

@@ -1,4 +1,5 @@
 import logging
+import sys
 
 logging.basicConfig(format="%(asctime)s %(name)s %(levelname)s %(message)s")
 log = logging.getLogger(__name__)
@@ -12,6 +13,8 @@
 
 import numpy as np
 
+from ..utils import KeepLoggingUnchanged
+
 __all__ = ["DataManagement"]
 
 
@@ -26,12 +29,13 @@ def findrun(run_number: int, search_on_GRID=True) -> Tuple[Path, List[Path]]:
         Returns:
             (PosixPath,list): the path list of *fits.fz files
         """
-        basepath = os.environ["NECTARCAMDATA"]
+        basepath = f"{os.environ['NECTARCAMDATA']}/runs/"
         list = glob.glob(
             basepath + "**/*" + str(run_number) + "*.fits.fz", recursive=True
         )
         list_path = [Path(chemin) for chemin in list]
         if len(list_path) == 0:
+            #############
             e = FileNotFoundError(f"run {run_number} is not present in {basepath}")
             if search_on_GRID:
                 log.warning(e, exc_info=True)
@@ -45,6 +49,7 @@ def findrun(run_number: int, search_on_GRID=True) -> Tuple[Path, List[Path]]:
             else:
                 log.error(e, exc_info=True)
                 raise e
+            ############## the pb is here !!!!!!
 
         name = list_path[0].name.split(".")
         name[2] = "*"
@@ -66,32 +71,86 @@ def getRunFromDIRAC(lfns: list):
         Args:
             lfns (list): list of lfns path
         """
-        from DIRAC.Interfaces.API.Dirac import Dirac
-
-        dirac = Dirac()
-        for lfn in lfns:
-            if not (
-                os.path.exists(f'{os.environ["NECTARCAMDATA"]}/{os.path.basename(lfn)}')
-            ):
-                dirac.getFile(
-                    lfn=lfn, destDir=os.environ["NECTARCAMDATA"], printOutput=True
-                )
+        with KeepLoggingUnchanged():
+            from DIRAC.Interfaces.API.Dirac import Dirac
+
+            dirac = Dirac()
+            for lfn in lfns:
+                if not (
+                    os.path.exists(
+                        f'{os.environ["NECTARCAMDATA"]}/runs/{os.path.basename(lfn)}'
+                    )
+                ):
+                    dirac.getFile(
+                        lfn=lfn,
+                        destDir=f"{os.environ['NECTARCAMDATA']}/runs/",
+                        printOutput=True,
+                    )
+                    pass
 
     @staticmethod
     def get_GRID_location(
-        run_number: int, output_lfns=True, username=None, password=None
+        run_number: int,
+        output_lfns=True,
+        basepath="/vo.cta.in2p3.fr/nectarcam/",
+        fromElog=False,
+        username=None,
+        password=None,
     ):
         """method to get run location on GRID from Elog (work in progress!)
 
         Args:
             run_number (int): run number
             output_lfns (bool, optional): if True, return lfns path of fits.gz files, else return parent directory of run location. Defaults to True.
+            basepath (str) : the path on GRID where nectarCAM data are stored. Default to /vo.cta.in2p3.fr/nectarcam/.
+            fromElog (bool,optionnl): To force to use the method which read the Elog. Default to False. To use the method with DIRAC API.
             username (_type_, optional): username for Elog login. Defaults to None.
             password (_type_, optional): password for Elog login. Defaults to None.
 
         Returns:
             _type_: _description_
         """
+        if fromElog:
+            return __class__.__get_GRID_location_ELog(
+                run_number=run_number,
+                output_lfns=output_lfns,
+                username=username,
+                password=password,
+            )
+        else:
+            return __class__.__get_GRID_location_DIRAC(
+                run_number=run_number, basepath=basepath
+            )
+
+    @staticmethod
+    def __get_GRID_location_DIRAC(
+        run_number: int, basepath="/vo.cta.in2p3.fr/nectarcam/"
+    ):
+        with KeepLoggingUnchanged():
+            from contextlib import redirect_stdout
+
+            from DIRAC.DataManagementSystem.Client.FileCatalogClientCLI import (
+                FileCatalogClientCLI,
+            )
+            from DIRAC.Interfaces.Utilities.DCommands import DCatalog
+
+            from nectarchain.utils import StdoutRecord
+
+            catalog = DCatalog()
+            with redirect_stdout(sys.stdout):
+                # stdout = StdoutRecord(keyword=f"Run{run_number}")
+                # fccli = FileCatalogClientCLI(catalog.catalog,stdout = stdout) // marche pas car DIRAC est mal code
+                fccli = FileCatalogClientCLI(catalog.catalog)
+                sys.stdout = StdoutRecord(keyword=f"Run{run_number}")
+                fccli.do_find("-q " + basepath)
+                lfns = sys.stdout.output
+                sys.stdout = sys.__stdout__
+        return lfns
+
+    @staticmethod
+    def __get_GRID_location_ELog(
+        run_number: int, output_lfns=True, username=None, password=None
+    ):
         import browser_cookie3
         import mechanize
         import requests

src/nectarchain/utils/__init__.py

@@ -1,2 +1,4 @@
 from .error import *
+from .io import *
+from .logger import *
 from .utils import *

src/nectarchain/utils/io.py

@@ -0,0 +1,22 @@
+import logging
+import sys
+
+logging.basicConfig(format="%(asctime)s %(name)s %(levelname)s %(message)s")
+log = logging.getLogger(__name__)
+log.handlers = logging.getLogger("__main__").handlers
+
+
+class StdoutRecord:
+    def __init__(self, keyword):
+        self.console = sys.stdout
+        self.keyword = keyword
+        self.output = []
+
+    def write(self, message):
+        if self.keyword in message:
+            self.console.write(message)
+            self.console.write("\n")
+            self.output.append(message)
+
+    def flush(self):
+        self.console.flush()

src/nectarchain/utils/logger.py

@@ -0,0 +1,27 @@
+import logging
+
+logging.basicConfig(format="%(asctime)s %(name)s %(levelname)s %(message)s")
+log = logging.getLogger(__name__)
+log.handlers = logging.getLogger("__main__").handlers
+
+import copy
+
+
+class KeepLoggingUnchanged:
+    def __init__(self):
+        self._nameToLevel = None
+        self._levelToName = None
+        self._srcfile = None
+        # self._lock = None
+
+    def __enter__(self):
+        self._nameToLevel = copy.copy(logging._nameToLevel)
+        self._levelToName = copy.copy(logging._levelToName)
+        self._srcfile = copy.copy(logging._srcfile)
+        # self._lock = copy.copy(logging._lock)
+
+    def __exit__(self, type, value, traceback):
+        logging._levelToName = self._levelToName
+        logging._nameToLevel = self._nameToLevel
+        logging._srcfile = self._srcfile
+        # logging._lock = self._lock