Coverage for src / piccione / upload / on_figshare.py: 100%

1import argparse 

2import os 

3import hashlib 

4import json 

5import requests 

6import yaml 

7from requests.exceptions import HTTPError 

8from tqdm import tqdm 

9 

10# Endpoint base di Figshare 

11BASE_URL = "https://api.figshare.com/v2/account/articles" 

12CHUNK_SIZE = 1048576 

13 

14 

15def get_file_check_data(file_name): 

16 with open(file_name, "rb") as fin: 

17 md5 = hashlib.md5() 

18 size = 0 

19 data = fin.read(CHUNK_SIZE) # circa 10MB 

20 while data: 

21 size += len(data) 

22 md5.update(data) 

23 data = fin.read(CHUNK_SIZE) 

24 return md5.hexdigest(), size 

25 

26 

27def issue_request(method, url, token, data=None, binary=False): 

28 headers = {"Authorization": "token " + token} 

29 if data is not None and not binary: 

30 data = json.dumps(data) 

31 response = requests.request(method, url, headers=headers, data=data) 

32 try: 

33 response.raise_for_status() 

34 try: 

35 data = json.loads(response.content) 

36 except ValueError: 

37 data = response.content 

38 except HTTPError as error: 

39 print(f"Caught an HTTPError: {str(error)}") 

40 print("Body:\n", response.text) 

41 raise 

42 return data 

43 

44 

45def upload_parts(file_info, file_path, token): 

46 url = file_info["upload_url"] 

47 result = issue_request(method="GET", url=url, token=token) 

48 print(f"\nUploading {os.path.basename(file_path)}:") 

49 

50 # Calculate total size for progress bar 

51 total_size = sum( 

52 part["endOffset"] - part["startOffset"] + 1 for part in result["parts"] 

53 ) 

54 

55 with open(file_path, "rb") as fin: 

56 with tqdm( 

57 total=total_size, unit="B", unit_scale=True, unit_divisor=1024 

58 ) as pbar: 

59 for part in result["parts"]: 

60 chunk_size = part["endOffset"] - part["startOffset"] + 1 

61 upload_part(file_info, fin, part, token) 

62 pbar.update(chunk_size) 

63 

64 

65def upload_part(file_info, stream, part, token): 

66 udata = file_info.copy() 

67 udata.update(part) 

68 url = "{upload_url}/{partNo}".format(**udata) 

69 stream.seek(part["startOffset"]) 

70 data = stream.read(part["endOffset"] - part["startOffset"] + 1) 

71 issue_request(method="PUT", url=url, data=data, binary=True, token=token) 

72 print(" Uploaded part {partNo} from {startOffset} to {endOffset}".format(**part)) 

73 

74 

75def create_file(article_id, file_name, file_path, token): 

76 url = f"{BASE_URL}/{article_id}/files" 

77 headers = {"Authorization": f"token {token}"} 

78 md5, size = get_file_check_data(file_path) 

79 data = {"name": os.path.basename(file_name), "md5": md5, "size": size} 

80 post_response = requests.post(url, headers=headers, json=data) 

81 get_response = requests.get(post_response.json()["location"], headers=headers) 

82 return get_response.json() 

83 

84 

85def complete_upload(article_id, file_id, token): 

86 url = f"{BASE_URL}/{article_id}/files/{file_id}" 

87 issue_request(method="POST", url=url, token=token) 

88 print(f" Upload completion confirmed for file {file_id}") 

89 

90 

91def main(config_path): 

92 with open(config_path) as f: 

93 config = yaml.safe_load(f) 

94 

95 token = config["TOKEN"] 

96 article_id = config["ARTICLE_ID"] 

97 files_to_upload = config["files_to_upload"] 

98 

99 print(f"Starting upload of {len(files_to_upload)} files to Figshare...") 

100 

101 for file_path in tqdm(files_to_upload, desc="Total progress", unit="file"): 

102 file_name = os.path.basename(file_path) 

103 print(f"\nPreparing {file_name}...") 

104 file_info = create_file(article_id, file_name, file_path, token) 

105 upload_parts(file_info, file_path, token) 

106 complete_upload(article_id, file_info["id"], token) 

107 print(f"✓ {file_name} completed") 

108 

109 print("\nAll files uploaded successfully to Figshare! 🎉") 

110 

111 

112if __name__ == "__main__": # pragma: no cover 

113 parser = argparse.ArgumentParser(description="Upload files to Figshare.") 

114 parser.add_argument("config", help="Path to the YAML configuration file.") 

115 args = parser.parse_args() 

116 main(args.config)