Working toward reading source from eggs, but this isn't right on Py3k yet.

8 files changed, 45 insertions, 7 deletions

diff --git a/Makefile b/Makefile
index c8e53c4e..71beee5b 100644
--- a/Makefile
+++ b/Makefile
@@ -4,6 +4,7 @@ default:
 	@echo "* No default action *"
 
 TEST_ZIP = test/zipmods.zip
+TEST_EGG = test/eggsrc/dist/covtestegg1-0.0.0-py2.5.egg
 
 clean:
 	python test/test_farm.py clean
@@ -16,6 +17,7 @@ clean:
 	-rm -f MANIFEST
 	-rm -f .coverage .coverage.* coverage.xml
 	-rm -f $(TEST_ZIP)
+	-rm -rf test/eggsrc/build test/eggsrc/dist test/eggsrc/*.egg-info
 	-rm -f setuptools-*.egg
 	-rm -rf doc/_build/*
 
@@ -44,10 +46,13 @@ testready: testdata devinst
 tests: testready
 	nosetests
 
-testdata: $(TEST_ZIP)
+testdata: $(TEST_ZIP) $(TEST_EGG)
 $(TEST_ZIP): test/covmodzip1.py
 	zip -j $@ $+
 
+$(TEST_EGG): test/eggsrc/setup.py test/eggsrc/egg1/egg1.py
+	cd test/eggsrc; python setup.py bdist_egg
+
 kit:
 	python setup.py sdist --formats=gztar
 	python setup.py bdist_wininst
diff --git a/coverage/backward.py b/coverage/backward.py
index d21e6a87..5e6b76fb 100644
--- a/coverage/backward.py
+++ b/coverage/backward.py
@@ -29,8 +29,9 @@ except NameError:
 
 try:
     from cStringIO import StringIO
+    BytesIO = StringIO
 except ImportError:
-    from io import StringIO
+    from io import StringIO, BytesIO
 
 # What's a string called?
 
diff --git a/coverage/codeunit.py b/coverage/codeunit.py
index 53c98fc7..3b0407ad 100644
--- a/coverage/codeunit.py
+++ b/coverage/codeunit.py
@@ -2,7 +2,8 @@
 
 import glob, os
 
-from coverage.backward import string_class
+from coverage.backward import string_class, BytesIO
+from coverage.misc import CoverageException
 
 
 def code_unit_factory(morfs, file_locator, omit_prefixes=None):
@@ -59,6 +60,8 @@ class CodeUnit:
     """
 
     def __init__(self, morf, file_locator):
+        self.file_locator = file_locator
+
         if hasattr(morf, '__file__'):
             f = morf.__file__
         else:
@@ -66,14 +69,14 @@ class CodeUnit:
         # .pyc files should always refer to a .py instead.
         if f.endswith('.pyc'):
             f = f[:-1]
-        self.filename = file_locator.canonical_filename(f)
+        self.filename = self.file_locator.canonical_filename(f)
 
         if hasattr(morf, '__name__'):
             n = modname = morf.__name__
             self.relative = True
         else:
             n = os.path.splitext(morf)[0]
-            rel = file_locator.relative_filename(n)
+            rel = self.file_locator.relative_filename(n)
             if os.path.isabs(n):
                 self.relative = (rel != n)
             else:
@@ -83,6 +86,7 @@ class CodeUnit:
         self.name = n
         self.modname = modname
 
+
     def __repr__(self):
         return "<CodeUnit name=%r filename=%r>" % (self.name, self.filename)
 
@@ -125,4 +129,12 @@ 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
+                    )
+            return BytesIO(source)
+        
         return open(self.filename)
diff --git a/coverage/files.py b/coverage/files.py
index c16f113d..cce33c97 100644
--- a/coverage/files.py
+++ b/coverage/files.py
@@ -34,7 +34,7 @@ class FileLocator:
         if filename not in self.canonical_filename_cache:
             f = filename
             if os.path.isabs(f) and not os.path.exists(f):
-                if not self.get_zip_data(f):
+                if self.get_zip_data(f) is None:
                     f = os.path.basename(f)
             if not os.path.isabs(f):
                 for path in [os.curdir] + sys.path:
@@ -50,7 +50,8 @@ class FileLocator:
         """Get data from `filename` if it is a zip file path.
         
         Returns the data read from the zip file, or None if no zip file could
-        be found or `filename` isn't in it.
+        be found or `filename` isn't in it.  The data returned might be "" if
+        the file is empty.
         
         """
         import zipimport
diff --git a/test/eggsrc/egg1/__init__.py b/test/eggsrc/egg1/__init__.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/test/eggsrc/egg1/__init__.py
diff --git a/test/eggsrc/egg1/egg1.py b/test/eggsrc/egg1/egg1.py
new file mode 100644
index 00000000..3fadde33
--- /dev/null
+++ b/test/eggsrc/egg1/egg1.py
@@ -0,0 +1,4 @@
+# My egg file!
+
+walrus = "Eggman"
+says = "coo-coo cachoo"
diff --git a/test/eggsrc/setup.py b/test/eggsrc/setup.py
new file mode 100644
index 00000000..6a88a58b
--- /dev/null
+++ b/test/eggsrc/setup.py
@@ -0,0 +1,6 @@
+from setuptools import setup
+
+setup(
+    name="covtestegg1",
+    packages=['egg1'],
+    )
diff --git a/test/test_codeunit.py b/test/test_codeunit.py
index 36ca1fcd..0c2fc98a 100644
--- a/test/test_codeunit.py
+++ b/test/test_codeunit.py
@@ -71,3 +71,12 @@ class CodeUnitTest(CoverageTest):
         assert zcu > acu and zcu >= acu and zcu != acu
         assert acu < bcu and acu <= bcu and acu != bcu
         assert bcu > acu and bcu >= acu and bcu != acu
+
+    def test_egg(self):
+        import egg1, egg1.egg1
+        cu = code_unit_factory([egg1, egg1.egg1], FileLocator())
+        self.assertEqual(cu[0].source_file().read(), "")
+        self.assertEqual(cu[1].source_file().read().split("\n")[0],
+                "# My egg file!"
+                )
+        
\ No newline at end of file