From 1525108e4b21c2617e839306d9ef6f15db18b9a6 Mon Sep 17 00:00:00 2001
From: David Sherret <dsherret@users.noreply.github.com>
Date: Sun, 5 Nov 2023 16:31:01 -0500
Subject: [PATCH] feat: ability to easily get types module from graph (#321)

---
 src/graph.rs | 22 ++++++++++++++++++++++
 src/lib.rs   | 23 +++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/src/graph.rs b/src/graph.rs
index 104eed6ab..dfed80955 100644
--- a/src/graph.rs
+++ b/src/graph.rs
@@ -1559,6 +1559,28 @@ impl ModuleGraph {
     }
   }
 
+  /// Similar to `try_get`, but will prefer resolving to the types dependency if
+  /// the module has one.
+  pub fn try_get_prefer_types(
+    &self,
+    specifier: &ModuleSpecifier,
+  ) -> Result<Option<&Module>, &ModuleError> {
+    let Some(module) = self.try_get(specifier)? else {
+      return Ok(None);
+    };
+
+    if let Some(specifier) = module.esm().and_then(|m| {
+      m.maybe_types_dependency
+        .as_ref()
+        .and_then(|d| d.dependency.ok())
+        .map(|r| &r.specifier)
+    }) {
+      self.try_get(specifier)
+    } else {
+      Ok(Some(module))
+    }
+  }
+
   /// Walk the graph from the root, checking to see if there are any module
   /// graph errors on non-type only, non-dynamic imports. The first error is
   /// returned as as error result, otherwise ok if there are no errors.
diff --git a/src/lib.rs b/src/lib.rs
index 4627ac6df..e20cd5f82 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -921,6 +921,29 @@ console.log(a);
         ModuleSpecifier::parse("https://example.com/jsx-runtime.d.ts").unwrap()
       )
     );
+    assert_eq!(
+      graph
+        .try_get(
+          &ModuleSpecifier::parse("https://example.com/jsx-runtime").unwrap()
+        )
+        .unwrap()
+        .unwrap()
+        .specifier()
+        .as_str(),
+      "https://example.com/jsx-runtime"
+    );
+    assert_eq!(
+      graph
+        .try_get_prefer_types(
+          &ModuleSpecifier::parse("https://example.com/jsx-runtime").unwrap()
+        )
+        .unwrap()
+        .unwrap()
+        .specifier()
+        .as_str(),
+      // should end up at the declaration file
+      "https://example.com/jsx-runtime.d.ts"
+    );
   }
 
   #[tokio::test]