Coverage for src/io_collection/save/save_text.py: 100%

1import io 

2from pathlib import Path 

3 

4from io_collection.save.save_buffer import _save_buffer_to_s3 

5 

6 

7def save_text( 

8 location: str, key: str, text: str, content_type: str = "binary/octet-stream" 

9) -> None: 

10 """ 

11 Save text to key at specified location. 

12 

13 Method will save to the S3 bucket if the location begins with the **s3://** 

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

15 

16 Parameters 

17 ---------- 

18 location 

19 Object location (local path or S3 bucket). 

20 key 

21 Object key ending in valid text extension. 

22 text 

23 Text to save. 

24 content_type 

25 Content type (S3 only). 

26 """ 

27 

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

29 _save_text_to_s3(location[5:], key, text, content_type) 

30 else: 

31 _save_text_to_fs(location, key, text) 

32 

33 

34def _save_text_to_fs(path: str, key: str, text: str) -> None: 

35 """ 

36 Save text to key on local file system. 

37 

38 Parameters 

39 ---------- 

40 path 

41 Local object path. 

42 key 

43 Object key ending in valid text extension. 

44 text 

45 Text to save. 

46 """ 

47 

48 full_path = Path(path) / key 

49 full_path.parent.mkdir(parents=True, exist_ok=True) 

50 with full_path.open("w") as fileobj: 

51 fileobj.write(text) 

52 

53 

54def _save_text_to_s3(bucket: str, key: str, text: str, content_type: str) -> None: 

55 """ 

56 Save text to key in AWS S3 bucket. 

57 

58 Parameters 

59 ---------- 

60 bucket 

61 AWS S3 bucket name. 

62 key 

63 Object key ending in valid text extension. 

64 text 

65 Text to save. 

66 content_type 

67 Content type (S3 only). 

68 """ 

69 

70 with io.BytesIO() as buffer: 

71 buffer.write(text.encode("utf-8")) 

72 _save_buffer_to_s3(bucket, key, buffer, content_type)