feat: Add a native dep parser request type

src/python/pants/backend/python/dependency_inference/parse_python_dependencies.py

@@ -19,6 +19,7 @@
 from pants.engine.environment import EnvironmentName
 from pants.engine.fs import CreateDigest, Digest, FileContent, MergeDigests
 from pants.engine.internals.native_dep_inference import NativeParsedPythonDependencies
+from pants.engine.internals.native_engine import NativeDependenciesRequest
 from pants.engine.process import Process, ProcessResult
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.unions import UnionMembership, UnionRule, union
@@ -195,7 +196,8 @@ async def parse_python_dependencies(
     has_custom_dep_inferences = len(union_membership[PythonDependencyVisitorRequest]) > 1
     if python_infer_subsystem.use_rust_parser and not has_custom_dep_inferences:
         native_result = await Get(
-            NativeParsedPythonDependencies, Digest, stripped_sources.snapshot.digest
+            NativeParsedPythonDependencies,
+            NativeDependenciesRequest(stripped_sources.snapshot.digest),
         )
         imports = dict(native_result.imports)
         assets = set()

src/python/pants/engine/internals/native_dep_inference.py

@@ -4,7 +4,10 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import Iterable
 
+from pants.engine.internals.native_engine import NativeDependenciesRequest
+from pants.engine.rules import QueryRule
 from pants.util.frozendict import FrozenDict
 
 
@@ -16,3 +19,8 @@ class NativeParsedPythonDependencies:
     def __init__(self, imports: dict[str, tuple[int, bool]], string_candidates: dict[str, int]):
         object.__setattr__(self, "imports", FrozenDict(imports))
         object.__setattr__(self, "string_candidates", FrozenDict(string_candidates))
+
+
+def rules() -> Iterable[QueryRule]:
+    # Keep in sync with `intrinsics.rs`.
+    return (QueryRule(NativeParsedPythonDependencies, (NativeDependenciesRequest,)),)

src/python/pants/engine/internals/native_engine.pyi

@@ -241,6 +241,28 @@ class PyStubCAS:
     def remove(self, digest: FileDigest | Digest) -> bool: ...
     def action_cache_len(self) -> int: ...
 
+# ------------------------------------------------------------------------------
+# Dependency inference
+# ------------------------------------------------------------------------------
+
+class NativeDependenciesRequest:
+    """A request to parse the dependencies of a file.
+
+    * The `digest` is expected to contain exactly one source file.
+    * Depending on the implementation, a json-serialized string `metadata`
+      can be passed. It will be supplied to the native parser, and
+      the string will be incorporated into the cache key.
+
+
+    Example:
+        result = await Get(NativeParsedPythonDependencies, NativeDependenciesRequest(input_digest, None)
+    """
+
+    def __init__(self, digest: Digest, metadata: str | None = None) -> None: ...
+    def __eq__(self, other: NativeDependenciesRequest | Any) -> bool: ...
+    def __hash__(self) -> int: ...
+    def __repr__(self) -> str: ...
+
 # ------------------------------------------------------------------------------
 # (etc.)
 # ------------------------------------------------------------------------------

src/python/pants/init/engine_initializer.py

@@ -25,6 +25,7 @@
     build_files,
     dep_rules,
     graph,
+    native_dep_inference,
     options_parsing,
     platform_rules,
     specs_rules,
@@ -278,6 +279,7 @@ def current_executing_goals(session_values: SessionValues) -> CurrentExecutingGo
                 *collect_rules(locals()),
                 *build_files.rules(),
                 *fs.rules(),
+                *native_dep_inference.rules(),
                 *dep_rules.rules(),
                 *desktop.rules(),
                 *download_file.rules(),

src/rust/engine/src/externs/dep_inference.rs

@@ -0,0 +1,60 @@
+// Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
+// Licensed under the Apache License, Version 2.0 (see LICENSE).
+use std::collections::hash_map::DefaultHasher;
+use std::hash::{Hash, Hasher};
+
+use pyo3::basic::CompareOp;
+use pyo3::prelude::*;
+use pyo3::{IntoPy, PyObject, Python};
+
+use fs::DirectoryDigest;
+
+use crate::externs::fs::PyDigest;
+
+pub(crate) fn register(m: &PyModule) -> PyResult<()> {
+  m.add_class::<PyNativeDependenciesRequest>()
+}
+
+#[pyclass(name = "NativeDependenciesRequest")]
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct PyNativeDependenciesRequest {
+  pub digest: DirectoryDigest,
+  pub metadata: Option<String>,
+}
+
+#[pymethods]
+impl PyNativeDependenciesRequest {
+  #[new]
+  fn __new__(digest: PyDigest, metadata: Option<String>) -> Self {
+    Self {
+      digest: digest.0,
+      metadata,
+    }
+  }
+
+  fn __hash__(&self) -> u64 {
+    let mut s = DefaultHasher::new();
+    self.digest.as_digest().hash.prefix_hash().hash(&mut s);
+    self.metadata.hash(&mut s);
+    s.finish()
+  }
+
+  fn __repr__(&self) -> String {
+    format!(
+      "NativeDependenciesRequest('{}', {})",
+      PyDigest(self.digest.clone()),
+      self
+        .metadata
+        .as_ref()
+        .map_or_else(|| "None", |string| string.as_str())
+    )
+  }
+
+  fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python) -> PyObject {
+    match op {
+      CompareOp::Eq => (self == other).into_py(py),
+      CompareOp::Ne => (self != other).into_py(py),
+      _ => py.NotImplemented(),
+    }
+  }
+}

src/rust/engine/src/externs/interface.rs

@@ -64,6 +64,7 @@ fn native_engine(py: Python, m: &PyModule) -> PyO3Result<()> {
   externs::scheduler::register(m)?;
   externs::testutil::register(m)?;
   externs::workunits::register(m)?;
+  externs::dep_inference::register(m)?;
 
   m.add("PollTimeout", py.get_type::<PollTimeout>())?;
 
@@ -234,6 +235,9 @@ impl PyTypes {
       docker_resolve_image_request: TypeId::new(docker_resolve_image_request),
       docker_resolve_image_result: TypeId::new(docker_resolve_image_result),
       parsed_python_deps_result: TypeId::new(parsed_python_deps_result),
+      deps_request: TypeId::new(
+        py.get_type::<externs::dep_inference::PyNativeDependenciesRequest>(),
+      ),
     })))
   }
 }

src/rust/engine/src/externs/mod.rs

@@ -23,6 +23,7 @@ use crate::interning::Interns;
 use crate::python::{Failure, Key, TypeId, Value};
 
 mod address;
+pub mod dep_inference;
 pub mod engine_aware;
 pub mod fs;
 mod interface;

src/rust/engine/src/intrinsics.rs

@@ -41,6 +41,7 @@ use rule_graph::DependencyKey;
 use stdio::TryCloneAsFile;
 use store::{SnapshotOps, SubsetParams};
 
+use crate::externs::dep_inference::PyNativeDependenciesRequest;
 use workunit_store::{in_workunit, Level};
 
 type IntrinsicFn =
@@ -50,7 +51,8 @@ pub struct Intrinsics {
   intrinsics: IndexMap<Intrinsic, IntrinsicFn>,
 }
 
-// NB: Keep in sync with `rules()` in `src/python/pants/engine/fs.py`.
+// NB: Keep in sync with `rules()` in `src/python/pants/engine/fs.py` and
+// `src/python/pants/engine/native_dep_inference.py`.
 impl Intrinsics {
   pub fn new(types: &Types) -> Intrinsics {
     let mut intrinsics: IndexMap<Intrinsic, IntrinsicFn> = IndexMap::new();
@@ -142,7 +144,7 @@ impl Intrinsics {
     intrinsics.insert(
       Intrinsic {
         product: types.parsed_python_deps_result,
-        inputs: vec![DependencyKey::new(types.directory_digest)],
+        inputs: vec![DependencyKey::new(types.deps_request)],
       },
       Box::new(parse_python_deps),
     );
@@ -754,10 +756,10 @@ fn parse_python_deps(context: Context, args: Vec<Value>) -> BoxFuture<'static, N
   async move {
     let core = &context.core;
     let store = core.store();
-    let directory_digest = Python::with_gil(|py| {
-      let py_digest = (*args[0]).as_ref(py);
-      lift_directory_digest(py_digest)
-    })?;
+    let PyNativeDependenciesRequest {
+      digest: directory_digest,
+      metadata,
+    } = Python::with_gil(|py| (*args[0]).as_ref(py).extract())?;
 
     let mut path = None;
     let mut digest = None;
@@ -786,7 +788,7 @@ fn parse_python_deps(context: Context, args: Vec<Value>) -> BoxFuture<'static, N
         let cache_key = CacheKey {
           key_type: CacheKeyType::DepInferenceRequest.into(),
           digest: Some(digest.into()),
-          metadata: None,
+          metadata: metadata,
         };
         let cached_result = core.local_cache.load(&cache_key).await?;
 

src/rust/engine/src/types.rs

@@ -35,4 +35,5 @@ pub struct Types {
   pub docker_resolve_image_request: TypeId,
   pub docker_resolve_image_result: TypeId,
   pub parsed_python_deps_result: TypeId,
+  pub deps_request: TypeId,
 }