Coverage for external_dependencies/utils.py: 24%

46 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-10-02 11:23 +0200

1import os 

2import zipfile 

3from importlib.metadata import distributions 

4 

5import boto3 

6from botocore import UNSIGNED 

7from botocore.config import Config 

8 

9s3_resource = boto3.resource('s3', config=Config(signature_version=UNSIGNED)) 

10 

11 

12def _list_external_dependencies() -> list: 

13 name_metadata_map = { 

14 dist.metadata['Name']: dist.metadata['Requires-External'] 

15 for dist in distributions() 

16 if 'Requires-External' in dist.metadata 

17 } 

18 requires_external = [] 

19 for _name, requires in name_metadata_map.items(): 

20 if type(requires) == str: 

21 requires_external.append(requires) 

22 elif type(requires) == list: 

23 requires_external.extend(requires) 

24 return list(set(requires_external)) 

25 

26 

27def _filter_external_dependencies(deps: list[str], prefix: str) -> list: 

28 return [dep[len(prefix) :] for dep in deps if dep.startswith(prefix)] 

29 

30 

31def _install_s3(dependencies: list[str], bucket_name: str, install_dir: str) -> None: 

32 if dependencies: 

33 print(f"Will install python 'Requires-External' for: {' '.join(dependencies)}") 

34 for dep in dependencies: 

35 print(f'Downloading {dep} from s3') 

36 bucket = s3_resource.Bucket(bucket_name) 

37 for obj in bucket.objects.filter(Prefix=f'{dep}/latest/'): 

38 if obj.key.endswith('/'): 

39 continue # Ignore directories 

40 filename = install_dir + obj.key.split('/')[-1] 

41 bucket.download_file(obj.key, filename) 

42 

43 with zipfile.ZipFile(filename, 'r') as zip_ref: 

44 for file in zip_ref.filelist: 

45 name = file.filename 

46 if name.endswith('/'): 

47 os.mkdir(os.path.join(install_dir, name)) 

48 os.chmod(os.path.join(install_dir, name), 0o755) 

49 else: 

50 outfile = os.path.join(install_dir, name) 

51 with open(outfile, 'bw') as file_fs: 

52 file_fs.write(zip_ref.read(name)) 

53 os.chmod(outfile, 0o755) 

54 os.remove(filename) 

55 else: 

56 print('No external dependencies found') 

57 

58 

59def _install_apt(dependencies: list[str]) -> None: 

60 if dependencies: 

61 print(f"Will install python 'Requires-External' for: {' '.join(dependencies)}") 

62 os.system(f'apt install -y {" ".join([dep for dep in dependencies])}') 

63 else: 

64 print('No external dependencies found')