Commits

Anonymous committed 2b515cf

now using dict and set as base type

Comments (0)

Files changed (2)

 #
 # DistributionDirectory represents a directory that contains egg-info files
 #
-class DistributionDirectory(object):
+class DistributionDirectory(set):
 
     def __init__(self, path):
+        super(DistributionDirectory, self).__init__()
         self.path = path
-        self._dists = []
         # filling the list once (see if it's the best way)
         # to minimize I/O
         for element in os.listdir(self.path):
             fullpath = join(self.path, element)
             if is_egg_info(fullpath):
-                self._dists.append(Distribution(fullpath))
+                self.add(Distribution(fullpath))
 
-    #
-    # container APIs
-    #
-    def __iter__(self):
-        return iter(self._dists)
+    def __repr__(self):
+        return 'DistributionDirectory("%s")' % self.path
+
+    def add(self, dist):
+        if not isinstance(dist, Distribution):
+            raise TypeError('DistributionDirectory manage only Distribution '
+                            'instances')
+        super(DistributionDirectory, self).add(dist)
 
     #
     # public APIs
     _CACHED_DIRS.clear()
 
 
-class DistributionDirectories(object):
+class DistributionDirectories(dict):
 
-    def __init__(self, use_cache=True):
-        self._paths = []
-        self._dist_dirs = []
+    def __init__(self, iterable=None, use_cache=True):
+        if iterable is None:
+            iterable = []
+        super(DistributionDirectories, self).__init__(iterable)
         self.use_cache = use_cache
 
+    def __setitem__(self, path, dir):
+        """Controls the mapping deals only with path/DistributionDirectory"""
+        if not isinstance(path, str):
+            raise TypeError('The key needs to be a path')
+        if not isinstance(dir, DistributionDirectory):
+            raise TypeError('The value needs to be a DistributionDirectory '
+                            'instance')
+        super(DistributionDirectories, self).__setitem__(path, dir)
+
     #
-    # container APIs
+    # public APIs
     #
-    def __iter__(self):
-        return iter(self._dist_dirs)
-
     def append(self, path):
+        """Appends a path."""
         if self.use_cache and path in _CACHED_DIRS:
             dist_dir = _CACHED_DIRS[path]
         else:
             if self.use_cache:
                 _CACHED_DIRS[path] = dist_dir
 
-        self._dist_dirs.append(dist_dir)
-        self._paths.append(path)
+        self[path] = dist_dir
 
-    def clear(self):
-        self._dist_dirs.clear()
-        self._paths.clear()
-
-    def remove(self, dist_dir):
-        self._dist_dirs.remove(dist_dir)
-        self._paths.remove(dist_dir.path)
-
-    #
-    # public APIs
-    #
     def load(self, paths):
         for path in paths:
-            if path in self._paths or not os.path.isdir(path):
+            if path in self or not os.path.isdir(path):
                 continue
             self.append(path)
 
     def reload(self):
-        paths = [d.path for d in self]
+        paths = self.keys()
         self.clear()
         self.load(paths)
 
     def get_distributions(self):
-        for directory in self:
+        for directory in self.values():
             for dist in directory:
                 yield dist
 
 
         If not found, returns None.
         """
-        for directory in self:
+        for directory in self.values():
             for dist in directory:
                 if dist.name == project_name:
                     return dist
 
         Return EggInfo instances.
         """
-        for directory in self:
+        for directory in self.values():
             for dist in directory.file_users(path):
                 yield dist
 
 def setup():
     sys.old = sys.path
     sys.path = [SITE_PKG]
+    pkgutil.purge_cache()
     pkgutil._dist_dirs = DistributionDirectories()
     pkgutil._dist_dirs.load([SITE_PKG])