Improve site package pathing

client/backend_arguments.py

@@ -69,7 +69,7 @@ def serialize(self) -> Dict[str, object]:
         }
 
     def get_checked_directory_allowlist(self) -> Set[str]:
-        return {element.path() for element in self.elements}
+        return {element.path() for element in self.elements if element is not None}
 
     def cleanup(self) -> None:
         pass
@@ -98,7 +98,7 @@ def serialize(self) -> Dict[str, object]:
         }
 
     def get_checked_directory_allowlist(self) -> Set[str]:
-        return {element.path() for element in self.elements}
+        return {element.path() for element in self.elements if element is not None}
 
     def cleanup(self) -> None:
         pass

client/configuration/search_path.py

@@ -19,13 +19,22 @@
 import glob
 import logging
 import os
-from typing import Dict, Iterable, List, Sequence, Union
+import re
+from typing import Dict, Iterable, List, Sequence, Tuple, Union
 
 from .. import filesystem
 from . import exceptions
 
 LOG: logging.Logger = logging.getLogger(__name__)
 
+dist_info_in_root: Dict[str, List[str]] = {}
+
+# pyre-fixme[5]: Globally accessible variable `_site_filter` has type `re.Pattern[str]` but no type is specified.
+_site_filter = re.compile(r".*-([0-99]\.)*dist-info")
+
+# pyre-fixme[5]: Globally accessible variable `_PYCACHE` has type `re.Pattern[str]` but no type is specified.
+_PYCACHE = re.compile("__pycache__(/)*.*")
+
 
 def _expand_relative_root(path: str, relative_root: str) -> str:
     if not path.startswith("//"):
@@ -35,11 +44,11 @@ def _expand_relative_root(path: str, relative_root: str) -> str:
 
 class Element(abc.ABC):
     @abc.abstractmethod
-    def path(self) -> str:
+    def path(self) -> Union[str, None]:
         raise NotImplementedError()
 
     @abc.abstractmethod
-    def command_line_argument(self) -> str:
+    def command_line_argument(self) -> Union[str, None]:
         raise NotImplementedError()
 
 
@@ -72,15 +81,55 @@ class SitePackageElement(Element):
     package_name: str
     is_toplevel_module: bool = False
 
-    def package_path(self) -> str:
-        module_suffix = ".py" if self.is_toplevel_module else ""
-        return self.package_name + module_suffix
+    def package_path(self) -> Union[str, None]:
+        if not self.is_toplevel_module:
+            return self.package_name
 
-    def path(self) -> str:
-        return os.path.join(self.site_root, self.package_path())
+        this_pkg_filter = re.compile(
+            r"{}-([0-99]\.)*dist-info(/)*.*".format(self.package_name)
+        )
 
-    def command_line_argument(self) -> str:
-        return self.site_root + "$" + self.package_path()
+        if self.site_root not in dist_info_in_root:
+            dist_info_in_root[self.site_root] = []
+            for directory in os.listdir(self.site_root):
+                if _site_filter.fullmatch(directory) is not None:
+                    dist_info_in_root[self.site_root].append(directory)
+
+        for directory in dist_info_in_root[self.site_root]:
+            if this_pkg_filter.fullmatch(directory):
+                dist_info_path = f"{self.site_root}/{directory}"
+                break
+        else:
+            return None
+
+        not_toplevel_patterns: Tuple[re.Pattern[str], re.Pattern[str]] = (
+            this_pkg_filter,
+            _PYCACHE,
+        )
+
+        # pyre-fixme[61]: Local variable `dist_info_path` is undefined, or not always defined.
+        with open(file=f"{dist_info_path}/RECORD", mode="r") as record:
+            for line in record.readlines():
+                file_name = line.split(",")[0]
+                for pattern in not_toplevel_patterns:
+                    if pattern.fullmatch(file_name):
+                        break
+                else:
+                    return file_name
+
+        return None
+
+    def path(self) -> Union[str, None]:
+        excepted_package_path: Union[str, None] = self.package_path()
+        if excepted_package_path is None:
+            return None
+        return os.path.join(self.site_root, excepted_package_path)
+
+    def command_line_argument(self) -> Union[str, None]:
+        excepted_package_path: Union[str, None] = self.package_path()
+        if excepted_package_path is None:
+            return None
+        return self.site_root + "$" + excepted_package_path
 
 
 class RawElement(abc.ABC):
@@ -223,10 +272,11 @@ def process_raw_elements(
     elements: List[Element] = []
 
     def add_if_exists(element: Element) -> bool:
-        if os.path.exists(element.path()):
-            elements.append(element)
-            return True
-        return False
+        excepted_path = element.path()
+        if excepted_path is None or not os.path.exists(excepted_path):
+            return False
+        elements.append(element)
+        return True
 
     for raw_element in raw_elements:
         expanded_raw_elements = raw_element.expand_glob()

client/configuration/tests/search_path_test.py

@@ -70,24 +70,45 @@ def test_create_raw_element(self) -> None:
             )
 
     def test_path(self) -> None:
-        self.assertEqual(SimpleElement("foo").path(), "foo")
-        self.assertEqual(SubdirectoryElement("foo", "bar").path(), "foo/bar")
-        self.assertEqual(SitePackageElement("foo", "bar").path(), "foo/bar")
+        with tempfile.TemporaryDirectory() as temp_root:
+            ensure_directories_exists(Path(temp_root), ("foo", "foo/bar"))
+            self.assertEqual(SimpleElement("foo").path(), "foo")
+            self.assertEqual(SubdirectoryElement("foo", "bar").path(), "foo/bar")
+            self.assertEqual(
+                SitePackageElement(f"{temp_root}/foo", "bar").path(),
+                f"{temp_root}/foo/bar",
+            )
 
     def test_command_line_argument(self) -> None:
-        self.assertEqual(SimpleElement("foo").command_line_argument(), "foo")
-        self.assertEqual(
-            SubdirectoryElement("foo", "bar").command_line_argument(),
-            "foo$bar",
-        )
-        self.assertEqual(
-            SitePackageElement("foo", "bar").command_line_argument(),
-            "foo$bar",
-        )
-        self.assertEqual(
-            SitePackageElement("foo", "bar", True).command_line_argument(),
-            "foo$bar.py",
-        )
+        with tempfile.TemporaryDirectory() as temp_root:
+            ensure_directories_exists(Path(temp_root), ("foo", "foo/bar"))
+            self.assertEqual(SimpleElement("foo").command_line_argument(), "foo")
+            self.assertEqual(
+                SubdirectoryElement("foo", "bar").command_line_argument(),
+                "foo$bar",
+            )
+            self.assertEqual(
+                SitePackageElement(f"{temp_root}/foo", "bar").command_line_argument(),
+                f"{temp_root}/foo$bar",
+            )
+
+        with tempfile.TemporaryDirectory() as temp_root:
+            ensure_directories_exists(
+                Path(temp_root), ("foo", "foo/bar-1.0.0.dist-info")
+            )
+            Path.touch(Path(f"{temp_root}/foo/bar.py"))
+
+            with open(
+                f"{temp_root}/foo/bar-1.0.0.dist-info/RECORD", "w", encoding="UTF-8"
+            ) as f:
+                f.write("bar.py,,")
+
+            self.assertEqual(
+                SitePackageElement(
+                    f"{temp_root}/foo", "bar", True
+                ).command_line_argument(),
+                f"{temp_root}/foo$bar.py",
+            )
 
     def test_expand_global_root(self) -> None:
         self.assertEqual(
@@ -241,3 +262,25 @@ def test_process_required_raw_elements_site_package_nonexistence(self) -> None:
                 site_roots=[],
                 required=True,
             )
+
+    def test_toplevel_module_not_pyfile(self) -> None:
+        with tempfile.TemporaryDirectory() as temp_root:
+            ensure_directories_exists(
+                Path(temp_root), ("foo", "foo/bar-1.0.0.dist-info")
+            )
+            Path.touch(Path(f"{temp_root}/foo/bar.so"))
+
+            with open(
+                f"{temp_root}/foo/bar-1.0.0.dist-info/RECORD", "w", encoding="UTF-8"
+            ) as f:
+                f.write("bar.so,,")
+
+            self.assertEqual(
+                SitePackageElement(f"{temp_root}/foo", "bar", True).path(), "foo/bar.so"
+            )
+            self.assertEqual(
+                process_raw_elements(
+                    [SitePackageRawElement("bar", True)], [f"{temp_root}/foo"]
+                ),
+                [SitePackageElement(f"{temp_root}/foo", "bar", True)],
+            )

client/tests/backend_arguments_test.py

@@ -3,6 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import shutil
 import tempfile
 from pathlib import Path
 from typing import Iterable, Tuple
@@ -576,15 +577,22 @@ def test_get_source_path__confliciting_source_specified(self) -> None:
             )
 
     def test_get_checked_directory_for_simple_source_path(self) -> None:
+        Path.mkdir(Path("super"))
+        Path.mkdir(Path("super/slash"))
+
         element0 = search_path.SimpleElement("ozzie")
         element1 = search_path.SubdirectoryElement("diva", "flea")
         element2 = search_path.SitePackageElement("super", "slash")
-        self.assertCountEqual(
-            SimpleSourcePath(
-                [element0, element1, element2, element0]
-            ).get_checked_directory_allowlist(),
-            [element0.path(), element1.path(), element2.path()],
-        )
+
+        try:
+            self.assertCountEqual(
+                SimpleSourcePath(
+                    [element0, element1, element2, element0]
+                ).get_checked_directory_allowlist(),
+                [element0.path(), element1.path(), element2.path()],
+            )
+        finally:
+            shutil.rmtree("super")
 
     def test_get_checked_directory_for_buck_source_path(self) -> None:
         self.assertCountEqual(