Coverage for src/io_collection/load/load_text.py: 100%

1from pathlib import Path 

2 

3from io_collection.load.load_buffer import _load_buffer_from_s3 

4 

5 

6def load_text(location: str, key: str) -> str: 

7 """ 

8 Load key as string from specified location. 

9 

10 Method will load from the S3 bucket if the location begins with the 

11 **s3://** 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 ending in valid text extension. 

19 

20 Returns 

21 ------- 

22 : 

23 Loaded text. 

24 """ 

25 

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

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

28 return _load_text_from_fs(location, key) 

29 

30 

31def _load_text_from_fs(path: str, key: str) -> str: 

32 """ 

33 Load key as dict or list from local file system. 

34 

35 Parameters 

36 ---------- 

37 path 

38 Local object path. 

39 key 

40 Object key ending in valid text extension. 

41 

42 Returns 

43 ------- 

44 : 

45 Loaded text. 

46 """ 

47 

48 full_path = Path(path) / key 

49 with full_path.open("r") as fileobj: 

50 return fileobj.read() 

51 

52 

53def _load_text_from_s3(bucket: str, key: str) -> str: 

54 """ 

55 Load key as dict or list from AWS S3 bucket. 

56 

57 Parameters 

58 ---------- 

59 bucket 

60 AWS S3 bucket name. 

61 key 

62 Object key ending in valid text extension. 

63 

64 Returns 

65 ------- 

66 : 

67 Loaded text. 

68 """ 

69 

70 buffer = _load_buffer_from_s3(bucket, key) 

71 return buffer.getvalue().decode("utf-8")