Coverage for cuda / pathfinder / _utils / find_sub_dirs.py: 95%

1# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# SPDX-License-Identifier: Apache-2.0 

3 

4import functools 

5import os 

6import site 

7import sys 

8from collections.abc import Sequence 

9 

10 

11def find_sub_dirs_no_cache(parent_dirs: Sequence[str], sub_dirs: Sequence[str]) -> list[str]: 

12 results = [] 

13 for base in parent_dirs: 

14 stack = [(base, 0)] # (current_path, index into sub_dirs) 

15 while stack: 

16 current_path, idx = stack.pop() 

17 if idx == len(sub_dirs): 

18 if os.path.isdir(current_path): 

19 results.append(current_path) 

20 continue 

21 

22 sub = sub_dirs[idx] 

23 if sub == "*": 

24 try: 

25 entries = sorted(os.listdir(current_path)) 

26 except OSError: 

27 continue 

28 for entry in entries: 

29 entry_path = os.path.join(current_path, entry) 

30 if os.path.isdir(entry_path): 

31 stack.append((entry_path, idx + 1)) 

32 else: 

33 next_path = os.path.join(current_path, sub) 

34 if os.path.isdir(next_path): 

35 stack.append((next_path, idx + 1)) 

36 return results 

37 

38 

39@functools.cache 

40def find_sub_dirs_cached(parent_dirs: Sequence[str], sub_dirs: Sequence[str]) -> list[str]: 

41 return find_sub_dirs_no_cache(parent_dirs, sub_dirs) 

42 

43 

44def find_sub_dirs(parent_dirs: Sequence[str], sub_dirs: Sequence[str]) -> list[str]: 

45 return find_sub_dirs_cached(tuple(parent_dirs), tuple(sub_dirs)) 

46 

47 

48def find_sub_dirs_sys_path(sub_dirs: Sequence[str]) -> list[str]: 

49 return find_sub_dirs(sys.path, sub_dirs) 

50 

51 

52def find_sub_dirs_all_sitepackages(sub_dirs: Sequence[str]) -> list[str]: 

53 return find_sub_dirs((site.getusersitepackages(), *site.getsitepackages()), sub_dirs)