Last clean up on the egg-reading code. Fixes issue #25.

3 files changed, 24 insertions, 10 deletions

diff --git a/CHANGES.txt b/CHANGES.txt
index b7dda4d5..bad54619 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -2,6 +2,14 @@
 Change history for Coverage.py

 ------------------------------

 

+Next version

+------------

+

+- Source code can now be read from eggs.  Thanks, rozza.  Fixes `issue 25`_

+

+.. _issue 25: http://bitbucket.org/ned/coveragepy/issue/25

+

+

 Version 3.1b1, 27 September 2009

 --------------------------------

 

diff --git a/TODO.txt b/TODO.txt
index 4cdb0514..64076924 100644
--- a/TODO.txt
+++ b/TODO.txt
@@ -2,7 +2,10 @@ Coverage TODO
 

 * 3.1

 

-- Python 3.1 support.

++ Python 3.1 support.

+

+* 3.2

+

 - bitvector in trace extension.

 

 

diff --git a/coverage/codeunit.py b/coverage/codeunit.py
index d85735de..79ae4f0e 100644
--- a/coverage/codeunit.py
+++ b/coverage/codeunit.py
@@ -86,7 +86,6 @@ class CodeUnit:
         self.name = n
         self.modname = modname
 
-
     def __repr__(self):
         return "<CodeUnit name=%r filename=%r>" % (self.name, self.filename)
 
@@ -129,12 +128,16 @@ class CodeUnit:
 
     def source_file(self):
         """Return an open file for reading the source of the code unit."""
-        if not os.path.exists(self.filename):
-            source = self.file_locator.get_zip_data(self.filename)
-            if source is None:
-                raise CoverageException(
-                    "No source for code %r." % self.filename
-                    )
+        if os.path.exists(self.filename):
+            # A regular text file: open it.
+            return open(self.filename)
+
+        # Maybe it's in a zip file?
+        source = self.file_locator.get_zip_data(self.filename)
+        if source is not None:
             return StringIO(source)
-        
-        return open(self.filename)
+            
+        # Couldn't find source.
+        raise CoverageException(
+            "No source for code %r." % self.filename
+            )