Implement pyre coverage (3/3) -- collect&print

Reviewed By: kbansal

Differential Revision: D30436430

fbshipit-source-id: b3f338c9bd2b604b568a0b176f30a2a2c15061c3

Author: grievejia

client/commands/v2/coverage.py

@@ -3,10 +3,21 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import dataclasses
+import json
 import logging
+import os
 from pathlib import Path
+from typing import Optional, Iterable, List
 
-from ... import commands, configuration as configuration_module
+import libcst as cst
+
+from ... import (
+    commands,
+    configuration as configuration_module,
+    coverage_collector as collector,
+    log,
+)
 from . import remote_logging, statistics
 
 LOG: logging.Logger = logging.getLogger(__name__)
@@ -22,13 +33,60 @@ def find_root(
     return working_directory
 
 
+@dataclasses.dataclass(frozen=True)
+class FileCoverage:
+    filepath: str
+    covered_lines: List[int]
+    uncovered_lines: List[int]
+
+
+def collect_coverage_for_module(relative_path: str, module: cst.Module) -> FileCoverage:
+    module_with_metadata = cst.MetadataWrapper(module)
+    coverage_collector = collector.CoverageCollector()
+    try:
+        module_with_metadata.visit(coverage_collector)
+    except RecursionError:
+        LOG.warning(f"LibCST encountered recursion error in `{relative_path}`")
+    return FileCoverage(
+        filepath=relative_path,
+        covered_lines=sorted(coverage_collector.covered_lines),
+        uncovered_lines=sorted(coverage_collector.uncovered_lines),
+    )
+
+
+def collect_coverage_for_path(
+    path: Path, working_directory: str
+) -> Optional[FileCoverage]:
+    module = statistics.parse_path_to_module(path)
+    relative_path = os.path.relpath(str(path), working_directory)
+    return (
+        collect_coverage_for_module(relative_path, module)
+        if module is not None
+        else None
+    )
+
+
+def collect_coverage_for_paths(
+    paths: Iterable[Path], working_directory: str
+) -> List[FileCoverage]:
+    result: List[FileCoverage] = []
+    for path in paths:
+        coverage = collect_coverage_for_path(path, working_directory)
+        if coverage is not None:
+            result.append(coverage)
+    return result
+
+
 def run_coverage(
     configuration: configuration_module.Configuration, working_directory: str
 ) -> commands.ExitCode:
-    sources = statistics.find_paths_to_parse(
-        [find_root(configuration, Path(working_directory))]
+    data = collect_coverage_for_paths(
+        statistics.find_paths_to_parse(
+            [find_root(configuration, Path(working_directory))]
+        ),
+        working_directory,
     )
-    LOG.warning(f"SOURCES = {list(sources)}")
+    log.stdout.write(json.dumps([dataclasses.asdict(entry) for entry in data]))
     return commands.ExitCode.SUCCESS
 
 

client/commands/v2/tests/coverage_test.py

@@ -3,12 +3,15 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import tempfile
 from pathlib import Path
+from typing import List
 
 import testslide
 
 from .... import configuration
-from ..coverage import find_root
+from ....tests import setup
+from ..coverage import find_root, collect_coverage_for_paths, FileCoverage
 
 
 class CoverageTest(testslide.TestCase):
@@ -33,3 +36,25 @@ def test_find_root(self) -> None:
             ),
             Path("/working/dir"),
         )
+
+    def test_collect_coverage(self) -> None:
+        with tempfile.TemporaryDirectory() as root:
+            root_path: Path = Path(root)
+            setup.ensure_files_exist(root_path, ["foo.py", "bar.py"])
+            foo_path = root_path / "foo.py"
+            bar_path = root_path / "bar.py"
+            baz_path = root_path / "baz.py"
+
+            data: List[FileCoverage] = collect_coverage_for_paths(
+                [foo_path, bar_path, baz_path], working_directory=root
+            )
+
+            def is_collected(path: Path) -> bool:
+                return any(
+                    str(path.relative_to(root_path)) == coverage.filepath
+                    for coverage in data
+                )
+
+            self.assertTrue(is_collected(foo_path))
+            self.assertTrue(is_collected(bar_path))
+            self.assertFalse(is_collected(baz_path))