Replace rdlogger.setLevel with supress_rdlogger decorator

Author: Oliver Scott

scaffoldgraph/core/fragment.py

@@ -27,8 +27,7 @@
 )
 
 from scaffoldgraph.core.scaffold import Scaffold
-
-rdlogger = RDLogger.logger()
+from scaffoldgraph.utils import supress_rdlogger
 
 
 class Fragmenter(ABC):
@@ -666,6 +665,7 @@ def get_annotated_murcko_scaffold(mol, scaffold=None, as_mol=True):
     return MolToSmiles(annotated)
 
 
+@supress_rdlogger()
 def get_next_murcko_fragments(murcko_scaffold, break_fused_rings=True):
     """
     Fragment a scaffold into its next set of Murcko fragments.
@@ -686,19 +686,16 @@ def get_next_murcko_fragments(murcko_scaffold, break_fused_rings=True):
         A list of parent scaffolds (next hierarchy [num_rings - 1])
 
     """
-    rdlogger.setLevel(4)
-
     if break_fused_rings:
         fragmenter = MurckoRingFragmenter()
     else:
         fragmenter = MurckoRingSystemFragmenter()
-
     parents = [f.mol for f in set(fragmenter.fragment(Scaffold(murcko_scaffold)))]
-    rdlogger.setLevel(3)
     return parents
 
 
 # fragmenter is hierarchical so all fragments may not be returned (fix?).
+@supress_rdlogger()
 def get_all_murcko_fragments(mol, break_fused_rings=True):
     """
     Get all possible murcko fragments from a molecule through
@@ -725,13 +722,10 @@ def get_all_murcko_fragments(mol, break_fused_rings=True):
     >>> frags = get_all_murcko_fragments(molecule)
 
     """
-    rdlogger.setLevel(4)
-
     if break_fused_rings:
         fragmenter = MurckoRingFragmenter()
     else:
         fragmenter = MurckoRingSystemFragmenter()
-
     mol = get_murcko_scaffold(mol)
     rdmolops.RemoveStereochemistry(mol)
     scaffold = Scaffold(mol)
@@ -745,7 +739,6 @@ def recursive_generation(child):
             recursive_generation(parent)
 
     recursive_generation(scaffold)
-    rdlogger.setLevel(3)
     return [f.mol for f in parents]
 
 

scaffoldgraph/core/graph.py

@@ -8,18 +8,18 @@
 from collections import Counter
 
 import networkx as nx
-import rdkit
 import gzip
 
 from loguru import logger
 from tqdm.auto import tqdm
 
 from rdkit import RDLogger
+from rdkit import __version__ as rdversion
 from rdkit.Chem import rdMolHash, MolToSmiles, rdmolops
 from rdkit.Chem.rdMolDescriptors import CalcNumRings
 
 from scaffoldgraph.io import *
-from scaffoldgraph.utils import canonize_smiles
+from scaffoldgraph.utils import canonize_smiles, supress_rdlogger
 
 from .fragment import (
     get_murcko_scaffold,
@@ -31,9 +31,6 @@
 
 from .scaffold import Scaffold
 
-rdlogger = RDLogger.logger()
-rdversion = rdkit.__version__
-
 
 def init_molecule_name(mol):
     """Initialize the name of a molecule if not provided.
@@ -149,6 +146,7 @@ def __init__(self, graph=None, fragmenter=None, graph_type=None, **attr):
         self.graph.setdefault('num_filtered', 0)  # track number of filtered mols.
         self.fragmenter = fragmenter
 
+    @supress_rdlogger()
     def _construct(self, molecules, init_args, ring_cutoff=10, progress=False):
         """Private method for graph construction, called by constructors.
 
@@ -179,7 +177,6 @@ def _construct(self, molecules, init_args, ring_cutoff=10, progress=False):
         _process_no_top_level
 
         """
-        rdlogger.setLevel(4)  # Supress rdkit logs.
         desc, progress = self.__class__.__name__, progress is False
         for molecule in tqdm(
             molecules, disable=progress, desc=desc,
@@ -196,7 +193,6 @@ def _construct(self, molecules, init_args, ring_cutoff=10, progress=False):
             scaffold = self._initialize_scaffold(molecule, init_args)
             if scaffold is not None:
                 self._hierarchy_constructor(scaffold)
-        rdlogger.setLevel(3)  # Enable rdkit logs.
 
     def _initialize_scaffold(self, molecule, init_args):
         """Initialize the top-level scaffold for a molecule.