Implement mapping of packages to instances of DependencyArg

in depgraph._get_arg_for_pkg(). Among other things, this
provides a way to know whether or not a specific package
has been specified as an argument, which determines whether
or not it can be removed from the graph during backtracking.

svn path=/main/trunk/; revision=8487

2 files changed, 50 insertions, 4 deletions

diff --git a/pym/_emerge/__init__.py b/pym/_emerge/__init__.py
index 3f79c420d..b6541bd47 100644
--- a/pym/_emerge/__init__.py
+++ b/pym/_emerge/__init__.py
@@ -951,12 +951,14 @@ class AtomArg(DependencyArg):
 	def __init__(self, atom=None, **kwargs):
 		DependencyArg.__init__(self, **kwargs)
 		self.atom = atom
+		self.set = (self.atom, )
 
 class PackageArg(DependencyArg):
 	def __init__(self, package=None, **kwargs):
 		DependencyArg.__init__(self, **kwargs)
 		self.package = package
 		self.atom = "=" + package.cpv
+		self.set = (self.atom, )
 
 class SetArg(DependencyArg):
 	def __init__(self, set=None, **kwargs):
@@ -1228,6 +1230,7 @@ class depgraph(object):
 		# contains all atoms from all sets added to the graph, including
 		# atoms given as arguments
 		self._set_atoms = InternalPackageSet()
+		self._atom_arg_map = {}
 		# contains all nodes pulled in by self._set_atoms
 		self._set_nodes = set()
 		self.blocker_digraph = digraph()
@@ -1700,9 +1703,15 @@ class depgraph(object):
 		# TODO: add multiple $ROOT support
 		if pkg.root != self.target_root:
 			return None
-		arg_atom = self._set_atoms.findAtomForPackage(pkg.cpv, pkg.metadata)
-		# TODO: map atom back to DependencyArg instance and return that instead
-		return arg_atom
+		atom_arg_map = self._atom_arg_map
+		any_arg = None
+		for atom in self._set_atoms.iterAtomsForPackage(pkg):
+			refs = atom_arg_map[(atom, pkg.root)]
+			for arg in refs:
+				any_arg = arg
+				if isinstance(arg, PackageArg):
+					return arg
+		return any_arg
 
 	def select_files(self, myfiles):
 		"""Given a list of .tbz2s, .ebuilds sets, and deps, create the
@@ -1886,6 +1895,16 @@ class depgraph(object):
 			if not oneshot:
 				myfavorites.append(myatom)
 		self._set_atoms.update(chain(*self._sets.itervalues()))
+		atom_arg_map = self._atom_arg_map
+		for arg in args:
+			for atom in arg.set:
+				atom_key = (atom, myroot)
+				refs = atom_arg_map.get(atom_key)
+				if refs is None:
+					refs = []
+					atom_arg_map[atom_key] = refs
+					if arg not in refs:
+						refs.append(arg)
 		pprovideddict = pkgsettings.pprovideddict
 		# Order needs to be preserved since a feature of --nodeps
 		# is to allow the user to force a specific merge order.
diff --git a/pym/portage/sets/base.py b/pym/portage/sets/base.py
index f360632ed..50702cff3 100644
--- a/pym/portage/sets/base.py
+++ b/pym/portage/sets/base.py
@@ -2,7 +2,7 @@
 # Distributed under the terms of the GNU General Public License v2
 # $Id$
 
-from portage import flatten
+from portage import cpv_getkey, flatten
 from portage.dep import isvalidatom, match_from_list, \
      best_match_to_list, dep_getkey, use_reduce, paren_reduce
 from portage.exception import InvalidAtom
@@ -126,6 +126,33 @@ class PackageSet(object):
 					return atoms[transformed_atoms.index(best_match)]
 		return None
 
+	def iterAtomsForPackage(self, pkg):
+		"""
+		Find all matching atoms for a given package. This matches virtual
+		arguments against the PROVIDE metadata.  This will raise an
+		InvalidDependString exception if PROVIDE is invalid.
+		"""
+		cpv_slot_list = ["%s:%s" % (pkg.cpv, pkg.metadata["SLOT"])]
+		cp = cpv_getkey(pkg.cpv)
+		self._load() # make sure the atoms are loaded
+		atoms = self._atommap.get(cp)
+		if atoms:
+			for atom in atoms:
+				if match_from_list(atom, cpv_slot_list):
+					yield atom
+		if not pkg.metadata["PROVIDE"]:
+			return
+		provides = flatten(use_reduce(paren_reduce(pkg.metadata["PROVIDE"]),
+			uselist=pkg.metadata["USE"].split()))
+		for provide in provides:
+			provided_cp = dep_getkey(provide)
+			atoms = self._atommap.get(provided_cp)
+			if atoms:
+				for atom in atoms:
+					if match_from_list(atom.replace(provided_cp, cp),
+						cpv_slot_list):
+						yield atom
+
 class EditablePackageSet(PackageSet):
 
 	def update(self, atoms):