Skip to content

Commit

Permalink
Fixed LD_LIBRARY_PATH squashing (#867)
Browse files Browse the repository at this point in the history
* Fixed LD_LIBRARY_PATH squashing

* Updated coverup test

* Added test of quoting everything

* Actually added the quoting
  • Loading branch information
sternj authored Oct 4, 2024
1 parent 4bb534f commit 856ba1a
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 16 deletions.
38 changes: 24 additions & 14 deletions scalene/scalene_preload.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,28 +34,38 @@ def get_preload_environ(args: argparse.Namespace, escape_spaces = False) -> Dict

elif sys.platform == "linux":
if args.memory:

library_path = scalene.__path__[0]


sanitized_path = scalene.__path__[0]
if escape_spaces:
# This function is used in two places, in `setup_preload` (where escaping spaces causes problems)
# and in `Scalene.__init__`. The latter creates a string by joining with spaces
# to pass into `redirect_python`, so we do need spaces there.
sanitized_path = sanitized_path.replace(" ", r"\ ")

# NOTE: you can't use escape sequences inside an f-string pre-3.12 either
if 'LD_LIBRARY_PATH' in env:
env['LD_LIBRARY_PATH'] = f'{sanitized_path}:{env["LD_LIBRARY_PATH"]}'
else:
env['LD_LIBRARY_PATH'] = sanitized_path

# We use this function in two places:
# 1. in `setup_preload`
# 2. when calling into `redirect_python`
if 'LD_LIBRARY_PATH' not in os.environ:

env['LD_LIBRARY_PATH'] = library_path
elif library_path not in os.environ['LD_LIBRARY_PATH']:
env['LD_LIBRARY_PATH'] = f'{library_path}:{os.environ["LD_LIBRARY_PATH"]}'


new_ld_preload = 'libscalene.so'
if "LD_PRELOAD" in env:
old_ld_preload = env["LD_PRELOAD"]
if "LD_PRELOAD" in os.environ and new_ld_preload not in os.environ["LD_PRELOAD"]:
old_ld_preload = os.environ["LD_PRELOAD"]
env["LD_PRELOAD"] = new_ld_preload + ":" + old_ld_preload
else:
env["LD_PRELOAD"] = new_ld_preload
# Disable command-line specified PYTHONMALLOC.
if "PYTHONMALLOC" in env:
del env["PYTHONMALLOC"]
if "PYTHONMALLOC" in os.environ:
# Since the environment dict is updated
# with a `.update` call, we need to make sure
# that there's some value for PYTHONMALLOC in
# what we return if we want to squash an anomalous
# value
env['PYTHONMALLOC'] = 'default'


elif sys.platform == "win32":
# Force no memory profiling on Windows for now.
Expand Down
2 changes: 1 addition & 1 deletion scalene/scalene_profiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -760,7 +760,7 @@ def __init__(
)
else:
preface = " ".join(
"=".join((k, str(v))) for (k, v) in environ.items()
"=".join((k, f"'{str(v)}'")) for (k, v) in environ.items()
)

Scalene.__orig_python = redirect_python(
Expand Down
2 changes: 1 addition & 1 deletion tests/test_coverup_5.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ def test_get_preload_environ_linux_memory(args, clean_environ):
env = scalene.scalene_preload.ScalenePreload.get_preload_environ(args)
assert 'LD_PRELOAD' in env
assert 'libscalene.so' in env['LD_PRELOAD']
assert 'PYTHONMALLOC' not in env
assert env['PYTHONMALLOC'] == 'default'

def test_get_preload_environ_linux_no_memory(args, clean_environ):

Expand Down

0 comments on commit 856ba1a

Please sign in to comment.