use std::sync::Arc;
use tokio::sync::mpsc::Sender;
use uv_distribution_types::{
Dist, DistributionId, Identifier, IndexMetadata, IndexUrl, Name, ResolvedDistRef,
};
use uv_normalize::PackageName;
use uv_once_map::Registration;
use uv_pep440::Version;
use crate::pubgrub::Range;
use crate::resolver::index::FxRegisteredEntry;
use crate::resolver::{InMemoryIndex, MetadataResponse, Request, VersionsResponse};
use crate::{PythonRequirement, ResolveError};
#[derive(Clone)]
pub(crate) struct MetadataRequests {
index: InMemoryIndex,
sender: Sender<Request>,
}
pub(crate) enum MetadataRequest<'a> {
Dist(Dist),
Resolved(ResolvedDistRef<'a>),
}
impl MetadataRequest<'_> {
fn id(&self) -> DistributionId {
match self {
Self::Dist(dist) => dist.distribution_id(),
Self::Resolved(dist) => dist.distribution_id(),
}
}
fn into_request(self) -> Request {
match self {
Self::Dist(dist) => Request::Dist(dist),
Self::Resolved(dist) => Request::from(dist),
}
}
}
impl Name for MetadataRequest<'_> {
fn name(&self) -> &PackageName {
match self {
Self::Dist(dist) => dist.name(),
Self::Resolved(dist) => dist.name(),
}
}
}
pub(crate) enum PendingVersions<'index> {
Implicit(FxRegisteredEntry<'index, PackageName, Arc<VersionsResponse>>),
Explicit(FxRegisteredEntry<'index, (PackageName, IndexUrl), Arc<VersionsResponse>>),
}
impl PendingVersions<'_> {
pub(crate) fn wait(self) -> Arc<VersionsResponse> {
match self {
Self::Implicit(entry) => entry.wait_blocking(),
Self::Explicit(entry) => entry.wait_blocking(),
}
}
}
#[derive(Clone, Debug)]
pub(crate) struct RegisteredMetadata<'index>(
FxRegisteredEntry<'index, DistributionId, Arc<MetadataResponse>>,
);
impl RegisteredMetadata<'_> {
pub(crate) fn id(&self) -> &DistributionId {
self.0.key()
}
pub(crate) fn wait(&self) -> Arc<MetadataResponse> {
self.0.wait_blocking()
}
}
impl MetadataRequests {
pub(crate) fn new(index: InMemoryIndex, sender: Sender<Request>) -> Self {
Self { index, sender }
}
pub(crate) fn enqueue_package(
&self,
name: &PackageName,
index: Option<&IndexMetadata>,
) -> Result<(), ResolveError> {
let registered = if let Some(index) = index {
self.index
.explicit()
.register((name.clone(), index.url().clone()))
} else {
self.index.implicit().register(name.clone())
};
if registered {
self.sender
.blocking_send(Request::Package(name.clone(), index.cloned()))?;
}
Ok(())
}
pub(crate) fn request_package(
&self,
name: &PackageName,
index: Option<&IndexMetadata>,
) -> Result<PendingVersions<'_>, ResolveError> {
if let Some(index) = index {
let entry = match self
.index
.explicit()
.register_entry((name.clone(), index.url().clone()))
{
Registration::New(entry) => {
self.sender
.blocking_send(Request::Package(name.clone(), Some(index.clone())))?;
entry
}
Registration::Existing(entry) => entry,
};
Ok(PendingVersions::Explicit(entry))
} else {
let entry = match self.index.implicit().register_entry(name.clone()) {
Registration::New(entry) => {
self.sender
.blocking_send(Request::Package(name.clone(), None))?;
entry
}
Registration::Existing(entry) => entry,
};
Ok(PendingVersions::Implicit(entry))
}
}
pub(crate) fn enqueue_metadata(
&self,
request: MetadataRequest<'_>,
) -> Result<(), ResolveError> {
if self.index.distributions().register(request.id()) {
self.sender.blocking_send(request.into_request())?;
}
Ok(())
}
pub(crate) fn request_metadata(
&self,
request: MetadataRequest<'_>,
validate: impl FnOnce(&MetadataRequest<'_>) -> Result<(), ResolveError>,
) -> Result<RegisteredMetadata<'_>, ResolveError> {
let entry = match self.index.distributions().register_entry(request.id()) {
Registration::New(entry) => {
validate(&request)?;
self.sender.blocking_send(request.into_request())?;
entry
}
Registration::Existing(entry) => entry,
};
Ok(RegisteredMetadata(entry))
}
pub(crate) fn prefetch(
&self,
name: &PackageName,
range: &Range<Version>,
python_requirement: &PythonRequirement,
) -> Result<(), ResolveError> {
self.sender.blocking_send(Request::Prefetch(
name.clone(),
range.clone(),
python_requirement.clone(),
))?;
Ok(())
}
pub(crate) fn metadata(&self, dist: &Dist) -> Result<RegisteredMetadata<'_>, ResolveError> {
self.index
.distributions()
.get_registered(dist.distribution_id())
.map(RegisteredMetadata)
.ok_or_else(|| ResolveError::UnregisteredTask(dist.to_string()))
}
}