bug: ship the modules json as a placeholder so uninstall removes it

lean/models/__init__.py

@@ -40,16 +40,16 @@ def _download_modules_json() -> None:
 
 error = None
 try:
-    # fetch if file not available or fetched before 1 day
-    if not path.exists(file_path) or (time() - path.getmtime(file_path) > 86400):
+    # fetch if file not available, still the empty placeholder the package ships, or fetched before 1 day
+    if not path.exists(file_path) or path.getsize(file_path) == 0 or (time() - path.getmtime(file_path) > 86400):
         _download_modules_json()
 except Exception as e:
     # No need to do anything if file isn't available
     error = str(e)
     pass
 
 # check if file exists
-if not Path(file_path).is_file():
+if not Path(file_path).is_file() or path.getsize(file_path) == 0:
     error_message = f": {error}" if error is not None else ""
     raise FileNotFoundError(
         f"Modules json not found in the given path {file_path}{error_message}")

setup.py

@@ -16,7 +16,9 @@
 import re
 import urllib.request
 
-from setuptools import find_packages, setup
+from setuptools import Command, find_packages, setup
+from setuptools.command.build_py import build_py
+from setuptools.command.sdist import sdist
 
 current_dir = os.path.abspath(os.path.dirname(__file__))
 
@@ -44,6 +46,42 @@ def get_stubs_version_range() -> str:
         return ""
 
 
+def write_modules_placeholder(command: Command, package_root: str) -> None:
+    """Writes an empty placeholder for the modules json the CLI downloads into its package, so pip records the file and
+    uninstalling removes it. Only build outputs get it, keeping a checkout's downloaded file intact."""
+    # A dry run changes nothing, and an in-place or linked build's package is the checkout's package itself
+    package_directory = os.path.join(package_root, "lean")
+    if command.dry_run or os.path.samefile(package_directory, os.path.join(current_dir, "lean")):
+        return
+
+    # The name is read rather than imported, as importing lean.models downloads the file
+    modules_file_name = re.search(r"^file_name = \"([^\"]+)\"", read("lean/models/__init__.py"), re.M).group(1)
+    placeholder_path = os.path.join(package_directory, modules_file_name)
+
+    # Build trees may hard link the checkout's files, so a copied file is replaced rather than truncated
+    if os.path.exists(placeholder_path):
+        os.remove(placeholder_path)
+    with open(placeholder_path, "w"):
+        pass
+
+
+class BuildPyWithModulesPlaceholder(build_py):
+    def run(self) -> None:
+        super().run()
+
+        # Editable builds link to the checkout, where the downloaded file has to stay
+        if getattr(self, "editable_mode", False):
+            return
+
+        write_modules_placeholder(self, self.build_lib)
+
+
+class SdistWithModulesPlaceholder(sdist):
+    def make_release_tree(self, base_dir: str, files: list) -> None:
+        super().make_release_tree(base_dir, files)
+        write_modules_placeholder(self, base_dir)
+
+
 # Production dependencies
 install_requires = [
     "click>=8.0.4",
@@ -75,6 +113,10 @@ def get_stubs_version_range() -> str:
     package_data={
         "lean": ["ssh/*"]
     },
+    cmdclass={
+        "build_py": BuildPyWithModulesPlaceholder,
+        "sdist": SdistWithModulesPlaceholder
+    },
     entry_points={
         "console_scripts": ["lean=lean.main:main"]
     },

tests/models/test_modules.py

@@ -11,9 +11,44 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import runpy
+from pathlib import Path
+
+import pytest
+from pyfakefs.fake_filesystem import FakeFilesystem
+from requests.exceptions import ConnectionError as RequestsConnectionError
+from responses import RequestsMock
+
+import lean.models
 from lean.models.modules import NuGetPackage
 
 
+def _write_modules_json_placeholder(fs: FakeFilesystem) -> None:
+    fs.add_real_file(lean.models.__file__)
+    fs.create_file(lean.models.file_path)
+
+
+def test_modules_json_placeholder_is_replaced_by_the_downloaded_file(fs: FakeFilesystem,
+                                                                      requests_mock: RequestsMock) -> None:
+    _write_modules_json_placeholder(fs)
+    modules = {"SomeModule": {"version": "1.0"}}
+    requests_mock.add(requests_mock.GET, lean.models.url, json={"modules": modules})
+
+    module_globals = runpy.run_path(lean.models.__file__)
+
+    assert module_globals["json_modules"] == modules
+    assert Path(lean.models.file_path).stat().st_size > 0
+
+
+def test_modules_json_placeholder_raises_not_found_when_download_fails(fs: FakeFilesystem,
+                                                                        requests_mock: RequestsMock) -> None:
+    _write_modules_json_placeholder(fs)
+    requests_mock.add(requests_mock.GET, lean.models.url, body=RequestsConnectionError())
+
+    with pytest.raises(FileNotFoundError, match="Modules json not found"):
+        runpy.run_path(lean.models.__file__)
+
+
 def test_get_file_name_returns_correct_value() -> None:
     package = NuGetPackage(name="QuantConnect.Lean", version="2.5.11940")
 

tests/test_setup.py

@@ -0,0 +1,111 @@
+# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
+# Lean CLI v1.0. Copyright 2021 QuantConnect Corporation.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import shutil
+import subprocess
+import sys
+import tarfile
+import tempfile
+import zipfile
+from pathlib import Path
+
+import pytest
+from pyfakefs.fake_filesystem import FakeFilesystem
+from setuptools.dist import Distribution
+
+import lean.models
+
+REPOSITORY_ROOT = Path(__file__).parent.parent
+MODULES_FILE = f"lean/{lean.models.file_name}"
+DOWNLOADED_MODULES_JSON = '{"modules": {}}'
+# Newer setuptools releases removed the --dry-run option
+requires_dry_run = pytest.mark.skipif("dry-run" not in [option[0] for option in Distribution.global_options],
+                                      reason="This setuptools version has no --dry-run option")
+
+
+@pytest.fixture
+def source_tree(fs: FakeFilesystem) -> Path:
+    """A minimal real copy of the package sources holding a downloaded modules json, as a checkout has."""
+    # setup.py runs in a subprocess, which needs the real filesystem
+    fs.pause()
+    with tempfile.TemporaryDirectory() as directory:
+        root = Path(directory)
+        for relative_path in ["setup.py", "README.md", "lean/__init__.py", "lean/models/__init__.py"]:
+            (root / relative_path).parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy(REPOSITORY_ROOT / relative_path, root / relative_path)
+        (root / MODULES_FILE).write_text(DOWNLOADED_MODULES_JSON)
+        yield root
+    fs.resume()
+
+
+def _run_setup(directory: Path, *arguments: str) -> None:
+    subprocess.run([sys.executable, "setup.py", *arguments], cwd=directory, check=True, capture_output=True)
+
+
+def _assert_wheel_has_empty_modules_placeholder(dist_directory: Path) -> None:
+    with zipfile.ZipFile(next(dist_directory.glob("*.whl"))) as wheel:
+        assert wheel.getinfo(MODULES_FILE).file_size == 0
+        assert MODULES_FILE in wheel.read(next(name for name in wheel.namelist() if name.endswith("RECORD"))).decode()
+
+
+def test_setup_metadata_keeps_downloaded_modules_json(source_tree: Path) -> None:
+    _run_setup(source_tree, "--version")
+
+    assert (source_tree / MODULES_FILE).read_text() == DOWNLOADED_MODULES_JSON
+
+
+@pytest.mark.parametrize("arguments", [
+    ["build_py", "--build-lib", "."],
+    pytest.param(["--dry-run", "build_py", "--build-lib", "."], marks=requires_dry_run),
+    pytest.param(["--dry-run", "build_py", "--build-lib", "build/lib"], marks=requires_dry_run)
+])
+def test_setup_in_place_and_dry_run_builds_keep_downloaded_modules_json(source_tree: Path, arguments: list) -> None:
+    _run_setup(source_tree, *arguments)
+
+    assert (source_tree / MODULES_FILE).read_text() == DOWNLOADED_MODULES_JSON
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Creating symbolic links may require elevated rights on Windows")
+def test_setup_build_into_linked_package_keeps_downloaded_modules_json(source_tree: Path) -> None:
+    (source_tree / "build" / "lib").mkdir(parents=True)
+    (source_tree / "build" / "lib" / "lean").symlink_to(source_tree / "lean", target_is_directory=True)
+
+    _run_setup(source_tree, "build_py")
+
+    assert (source_tree / MODULES_FILE).read_text() == DOWNLOADED_MODULES_JSON
+
+
+def test_setup_wheel_records_empty_modules_placeholder_and_keeps_downloaded_file(source_tree: Path) -> None:
+    _run_setup(source_tree, "bdist_wheel")
+
+    _assert_wheel_has_empty_modules_placeholder(source_tree / "dist")
+    assert (source_tree / MODULES_FILE).read_text() == DOWNLOADED_MODULES_JSON
+
+
+def test_setup_sdist_has_empty_modules_placeholder_and_keeps_downloaded_file(source_tree: Path) -> None:
+    _run_setup(source_tree, "sdist")
+
+    with tarfile.open(next((source_tree / "dist").glob("*.tar.gz"))) as sdist:
+        assert next(member for member in sdist.getmembers() if member.name.endswith(MODULES_FILE)).size == 0
+    assert (source_tree / MODULES_FILE).read_text() == DOWNLOADED_MODULES_JSON
+
+
+def test_setup_wheel_built_from_sdist_records_empty_modules_placeholder(source_tree: Path) -> None:
+    _run_setup(source_tree, "sdist")
+    with tarfile.open(next((source_tree / "dist").glob("*.tar.gz"))) as sdist:
+        sdist.extractall(source_tree / "extracted")
+    sdist_tree = next((source_tree / "extracted").iterdir())
+
+    _run_setup(sdist_tree, "bdist_wheel")
+
+    _assert_wheel_has_empty_modules_placeholder(sdist_tree / "dist")