Skip to content

Commit d9bbdbc

Browse files
soiferjjroesch
authored andcommitted
[PYTHON/FFI] Search PATH for DLLs (#3888)
* Search PATH for DLLs * Fix lint issue
1 parent 08d9220 commit d9bbdbc

File tree

1 file changed

+27
-4
lines changed

1 file changed

+27
-4
lines changed

python/tvm/_ffi/libinfo.py

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,25 @@
1919
import sys
2020
import os
2121

22+
def split_env_var(env_var, split):
23+
"""Splits environment variable string.
24+
25+
Parameters
26+
----------
27+
env_var : str
28+
Name of environment variable.
29+
30+
split : str
31+
String to split env_var on.
32+
33+
Returns
34+
-------
35+
splits : list(string)
36+
If env_var exists, split env_var. Otherwise, empty list.
37+
"""
38+
if os.environ.get(env_var, None):
39+
return [p.strip() for p in os.environ[env_var].split(split)]
40+
return []
2241

2342
def find_lib_path(name=None, search_path=None, optional=False):
2443
"""Find dynamic library files.
@@ -50,10 +69,14 @@ def find_lib_path(name=None, search_path=None, optional=False):
5069
if os.environ.get('TVM_LIBRARY_PATH', None):
5170
dll_path.append(os.environ['TVM_LIBRARY_PATH'])
5271

53-
if sys.platform.startswith('linux') and os.environ.get('LD_LIBRARY_PATH', None):
54-
dll_path.extend([p.strip() for p in os.environ['LD_LIBRARY_PATH'].split(":")])
55-
elif sys.platform.startswith('darwin') and os.environ.get('DYLD_LIBRARY_PATH', None):
56-
dll_path.extend([p.strip() for p in os.environ['DYLD_LIBRARY_PATH'].split(":")])
72+
if sys.platform.startswith('linux'):
73+
dll_path.extend(split_env_var('LD_LIBRARY_PATH', ':'))
74+
dll_path.extend(split_env_var('PATH', ':'))
75+
elif sys.platform.startswith('darwin'):
76+
dll_path.extend(split_env_var('DYLD_LIBRARY_PATH', ':'))
77+
dll_path.extend(split_env_var('PATH', ':'))
78+
elif sys.platform.startswith('win32'):
79+
dll_path.extend(split_env_var('PATH', ';'))
5780

5881
# Pip lib directory
5982
dll_path.append(os.path.join(ffi_dir, ".."))

0 commit comments

Comments
 (0)