diff --git a/google/colab/_pip.py b/google/colab/_pip.py
index 0525c5ab..41a05433 100644
--- a/google/colab/_pip.py
+++ b/google/colab/_pip.py
@@ -4,10 +4,10 @@
 that is already loaded in sys.modules.
 """
 
+import collections
+import importlib
 import itertools
-import os
 import re
-import site
 import sys
 import uuid
 from IPython.core.display import _display_mimetype
@@ -33,60 +33,16 @@ def _extract_installed_packages(pip_output):
   return itertools.chain(*map(str.split, results))
 
 
-def _get_distinfo_path(distname, paths):
-  """Find the filesystem path to a package's distribution info.
-
-  Distribution names must be treated as case-insensitive, with '-' and '_'
-  characters treated as equivalent
-  (See https://www.python.org/dev/peps/pep-0426/#name).
-
-  Args:
-    distname: distribution name.
-    paths: list of directory path to search
-
-  Returns:
-    path: (string or None) the valid filesystem path to the distribution.
-  """
-  paths = [p for p in paths if os.path.exists(p)]
-  if not paths:
-    return None
-
-  # Python packages can be installed as wheels or as eggs. Account for both
-  # (see https://packaging.python.org/discussions/wheel-vs-egg/)
-  distinfo = ["{}.dist-info".format(distname), "{}.egg-info".format(distname)]
-
-  def normalize_dist(dist):
-    return dist.lower().replace("_", "-")
-
-  distinfo = [normalize_dist(info) for info in distinfo]
-
-  for path in paths:
-    path_map = {normalize_dist(f): f for f in os.listdir(path)}
-    for info in distinfo:
-      if info in path_map:
-        joined = os.path.join(path, path_map[info])
-        if os.path.isdir(joined):
-          return joined
-
-  return None
-
-
 def _extract_toplevel_packages(pip_output):
   """Extract the list of toplevel packages associated with a pip install."""
-  # Account for default installations and --user installations (most common).
-  # Note: we should possibly also account for --root, --prefix, & -t/--target.
-  sitepackages = site.getsitepackages() + [site.getusersitepackages()]
-  for package in _extract_installed_packages(pip_output):
-    infodir = _get_distinfo_path(package, sitepackages)
-    if not infodir:
-      continue
-    toplevel = os.path.join(infodir, "top_level.txt")
-    if not os.path.exists(toplevel):
-      continue
-    for line in open(toplevel):
-      line = line.strip()
-      if line:
-        yield line
+  toplevel = collections.defaultdict(set)
+  for m, ps in importlib.metadata.packages_distributions().items():
+    for p in ps:
+      toplevel[p].add(m)
+  for pv in _extract_installed_packages(pip_output):
+    package = pv[: pv.rfind("-")]
+    for module in toplevel.get(package, set()):
+      yield module
 
 
 def _previously_imported_packages(pip_output):