Fix pkg_resources warning (#804)

Author: Logiquo

pyhealth/models/molerec.py

@@ -1,7 +1,8 @@
 import torch
 import os
 import math
-import pkg_resources
+import importlib.metadata as importlib_metadata
+from packaging import version
 import numpy as np
 from typing import Any, Dict, List, Tuple, Optional, Union
 from rdkit import Chem
@@ -18,6 +19,31 @@
 
 from pyhealth import BASE_CACHE_PATH as CACHE_PATH
 
+
+def _import_ogb(min_version: str = "1.3.5") -> None:
+    """Import ogb components and validate version."""
+    try:
+        installed = importlib_metadata.version("ogb")
+    except importlib_metadata.PackageNotFoundError as exc:
+        raise ImportError(f"Please install ogb>={min_version}.") from exc
+
+    if version.parse(installed) < version.parse(min_version):
+        raise ImportError(
+            f"ogb>={min_version} required, but found {installed}."
+        )
+
+    from ogb.utils import smiles2graph
+    from ogb.graphproppred.mol_encoder import AtomEncoder, BondEncoder
+
+    globals().update(
+        {
+            "smiles2graph": smiles2graph,
+            "AtomEncoder": AtomEncoder,
+            "BondEncoder": BondEncoder,
+        }
+    )
+
+
 def graph_batch_from_smiles(smiles_list, device=torch.device("cpu")):
     edge_idxes, edge_feats, node_feats, lstnode, batch = [], [], [], 0, []
 
@@ -309,19 +335,14 @@ def __init__(
     ):
         super(MoleRecLayer, self).__init__()
 
-        dependencies = ["ogb>=1.3.5"]
-
         # test whether the ogb and torch_scatter packages are ready
         try:
-            pkg_resources.require(dependencies)
-            global smiles2graph, AtomEncoder, BondEncoder
-            from ogb.utils import smiles2graph
-            from ogb.graphproppred.mol_encoder import AtomEncoder, BondEncoder
-        except Exception as e:
+            _import_ogb("1.3.5")
+        except Exception:
             print(
                 "Please follow the error message and install the [ogb>=1.3.5] packages first."
             )
-            print(e)
+            raise
 
         self.hidden_size = hidden_size
         self.coef, self.target_ddi = coef, target_ddi
@@ -569,18 +590,13 @@ def __init__(
             else:
                 self.mode = "multilabel"
 
-        dependencies = ["ogb>=1.3.5"]
-
         try:
-            pkg_resources.require(dependencies)
-            global smiles2graph, AtomEncoder, BondEncoder
-            from ogb.utils import smiles2graph
-            from ogb.graphproppred.mol_encoder import AtomEncoder, BondEncoder
-        except Exception as e:
+            _import_ogb("1.3.5")
+        except Exception:
             print(
                 "Please follow the error message and install the [ogb>=1.3.5] packages first."
             )
-            print(e)
+            raise
 
         self.embedding_dim = embedding_dim
         self.hidden_dim = hidden_dim

pyhealth/sampler/sage_sampler.py

@@ -1,6 +1,5 @@
 import copy
 from typing import List, Optional, Tuple, NamedTuple, Union, Callable, TypeVar
-import pkg_resources
 import torch
 from torch import Tensor
 
@@ -103,19 +102,16 @@ def __init__(self, dataset, edge_index, # is a Union[Tensor, SparseTensor],
 
         if 'collate_fn' in kwargs:
             del kwargs['collate_fn']
-            
-        dependencies = ["torch_sparse"]
 
-        # test whether the ogb and torch_scatter packages are ready
+        # Ensure torch_sparse is available before proceeding
         try:
-            pkg_resources.require(dependencies)
             global SparseTensor
             from torch_sparse import SparseTensor
         except Exception as e:
             print(
                 "Please follow the error message and install the [torch_sparse] packages first."
             )
-            print(e)
+            raise
 
         # Save for Pytorch Lightning...
         self.raw_data = dataset

tests/core/test_molerec.py

@@ -1,5 +1,7 @@
 import unittest
 import torch
+import importlib.metadata as importlib_metadata
+from packaging import version
 
 from pyhealth.datasets import create_sample_dataset, get_dataloader
 from pyhealth.models import MoleRec
@@ -10,11 +12,19 @@ class TestMoleRec(unittest.TestCase):
 
     def setUp(self):
         """Set up test data and model."""
+        min_version = "1.3.5"
         try:
-            import pkg_resources
-            pkg_resources.require(["ogb>=1.3.5"])
-        except Exception:
+            installed = importlib_metadata.version("ogb")
+        except importlib_metadata.PackageNotFoundError:
             self.skipTest("ogb package not available")
+
+        if version.parse(installed) < version.parse(min_version):
+            self.skipTest(f"ogb>={min_version} required, but found {installed}")
+
+        try:
+            from ogb.utils import smiles2graph  # noqa: F401
+        except Exception as exc:
+            self.skipTest(f"ogb import failed: {exc}")
 
         self.samples = [
             {