Look for __main__ module if coverage is being run for directory #252

4 files changed, 22 insertions, 0 deletions

diff --git a/coverage/execfile.py b/coverage/execfile.py
index 2d856897..71bdb8db 100644
--- a/coverage/execfile.py
+++ b/coverage/execfile.py
@@ -137,6 +137,16 @@ def run_python_file(filename, args, package=None, modulename=None):
     old_argv = sys.argv
     sys.argv = args
 
+    if os.path.isdir(filename):
+        # in directory we should look for __main__ module
+        for ext in [".py", ".pyc", ".pyo"]:
+            try_filename = os.path.join(filename, "__main__" + ext)
+            if os.path.exists(try_filename):
+                filename = try_filename
+                break
+        else:
+            raise NoSource("Can't find '__main__' module in '%s'" % filename)
+
     try:
         # Make a code object somehow.
         if filename.endswith((".pyc", ".pyo")):
diff --git a/tests/test_execfile.py b/tests/test_execfile.py
index 49d6cc8c..9c2bc7d6 100644
--- a/tests/test_execfile.py
+++ b/tests/test_execfile.py
@@ -82,6 +82,15 @@ class RunFileTest(CoverageTest):
         with self.assertRaises(NoSource):
             run_python_file("xyzzy.py", [])
 
+    def test_directory_with_main(self):
+        directory_with_main = os.path.join(HERE, "with_main")
+        run_python_file(directory_with_main, [directory_with_main])
+        self.assertEqual(self.stdout(), "1\n")
+
+    def test_directory_without_main(self):
+        with self.assertRaises(NoSource):
+            directory_with_main = os.path.join(HERE, "with_main", "without")
+            run_python_file(directory_with_main, [directory_with_main])
 
 class RunPycFileTest(CoverageTest):
     """Test cases for `run_python_file`."""
diff --git a/tests/with_main/__main__.py b/tests/with_main/__main__.py
new file mode 100644
index 00000000..e7a4e4f9
--- /dev/null
+++ b/tests/with_main/__main__.py
@@ -0,0 +1,2 @@
+x = 1
+print(x)
\ No newline at end of file
diff --git a/tests/with_main/without/__init__.py b/tests/with_main/without/__init__.py
new file mode 100644
index 00000000..595e3818
--- /dev/null
+++ b/tests/with_main/without/__init__.py
@@ -0,0 +1 @@
+__author__ = 'traff'