Coverage for src/io_collection/keys/check_key.py: 100%

1from pathlib import Path 

2 

3import boto3 

4 

5 

6def check_key(location: str, key: str) -> bool: 

7 """ 

8 Check if object key exists at specified location. 

9 

10 Method will check in the S3 bucket if the location begins with the **s3://** 

11 protocol, otherwise it assumes the location is a local path. 

12 

13 Parameters 

14 ---------- 

15 location 

16 Object location (local path or S3 bucket). 

17 key 

18 Object key. 

19 

20 Returns 

21 ------- 

22 : 

23 True if the object exists in the location, False otherwise. 

24 """ 

25 

26 if location[:5] == "s3://": 

27 return _check_key_on_s3(location[5:], key) 

28 return _check_key_on_fs(location, key) 

29 

30 

31def _check_key_on_fs(path: str, key: str) -> bool: 

32 """ 

33 Check if object key exists on local file system. 

34 

35 Parameters 

36 ---------- 

37 path 

38 Local object path. 

39 key 

40 Object key. 

41 

42 Returns 

43 ------- 

44 bool 

45 True if the object exists on the local file system, False otherwise. 

46 """ 

47 

48 full_path = Path(path) / key 

49 return full_path.is_file() 

50 

51 

52def _check_key_on_s3(bucket: str, key: str) -> bool: 

53 """ 

54 Check if object key exists in AWS S3 bucket. 

55 

56 Parameters 

57 ---------- 

58 bucket 

59 AWS S3 bucket name. 

60 key 

61 Object key. 

62 

63 Returns 

64 ------- 

65 bool 

66 True if the object exists in the AWS bucket, False otherwise. 

67 """ 

68 

69 s3_client = boto3.client("s3") 

70 result = s3_client.list_objects(Bucket=bucket, Prefix=key) 

71 return "Contents" in result