|
19 | 19 | import sys |
20 | 20 | import os |
21 | 21 |
|
| 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 [] |
22 | 41 |
|
23 | 42 | def find_lib_path(name=None, search_path=None, optional=False): |
24 | 43 | """Find dynamic library files. |
@@ -50,10 +69,14 @@ def find_lib_path(name=None, search_path=None, optional=False): |
50 | 69 | if os.environ.get('TVM_LIBRARY_PATH', None): |
51 | 70 | dll_path.append(os.environ['TVM_LIBRARY_PATH']) |
52 | 71 |
|
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', ';')) |
57 | 80 |
|
58 | 81 | # Pip lib directory |
59 | 82 | dll_path.append(os.path.join(ffi_dir, "..")) |
|
0 commit comments