Put caller cwd on sys.path for generate metadata

petastorm/etl/petastorm_generate_metadata.py

@@ -15,6 +15,7 @@
 """Script to add petastorm metadata to an existing parquet dataset"""
 
 import argparse
+import os
 import sys
 from pydoc import locate
 
@@ -27,6 +28,19 @@
 from petastorm.unischema import Unischema
 from petastorm.utils import add_to_dataset_metadata
 
+
+def _ensure_cwd_on_sys_path():
+    """Put the caller's working directory first on ``sys.path``.
+
+    Console-script installs run from the environment's ``bin`` directory, so
+    ``pydoc.locate`` cannot find a unischema module that lives on the user's
+    current python path / project tree unless cwd is searched first.
+    """
+    cwd = os.getcwd()
+    if sys.path[:1] != [cwd]:
+        sys.path.insert(0, cwd)
+
+
 example_text = '''Example (some replacement required):
 
 Locally:
@@ -112,6 +126,10 @@ def generate_petastorm_metadata(spark, dataset_url, unischema_class=None, use_su
 
 
 def _main(args):
+    # Console scripts may start with the venv bin dir as the process context;
+    # prefer the caller's cwd so --unischema_class can resolve project modules.
+    _ensure_cwd_on_sys_path()
+
     parser = argparse.ArgumentParser(prog='petastorm_generate_metadata',
                                      description='Add necessary petastorm metadata to an existing dataset',
                                      epilog=example_text,

petastorm/tests/test_generate_metadata.py

@@ -29,6 +29,19 @@ def _check_reader(path, rowgroup_selector=None):
         [next(reader) for _ in range(10)]
 
 
+def test_ensure_cwd_on_sys_path_puts_cwd_first(monkeypatch, tmpdir):
+    # Simulate a console-script style path where the install dir is first and
+    # the caller's project cwd is missing from sys.path.
+    project_cwd = tmpdir.mkdir('project').strpath
+    install_bin = tmpdir.mkdir('bin').strpath
+    monkeypatch.chdir(project_cwd)
+    monkeypatch.setattr(petastorm_generate_metadata.sys, 'path', [install_bin])
+
+    petastorm_generate_metadata._ensure_cwd_on_sys_path()
+
+    assert petastorm_generate_metadata.sys.path[0] == os.path.abspath(project_cwd)
+
+
 def test_regenerate_metadata(synthetic_dataset, tmpdir):
     a_moved_path = tmpdir.join('moved').strpath
     copytree(synthetic_dataset.path, a_moved_path)