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

1import argparse 

2import time 

3from pathlib import Path 

4 

5import requests 

6import yaml 

7from tqdm import tqdm 

8 

9 

10class ProgressFileWrapper: 

11 def __init__(self, file_path): 

12 self.file_path = file_path 

13 self.file_size = Path(file_path).stat().st_size 

14 self.fp = open(file_path, 'rb') 

15 self.pbar = tqdm(total=self.file_size, unit='B', unit_scale=True, desc=Path(file_path).name) 

16 

17 def read(self, size=-1): 

18 data = self.fp.read(size) 

19 self.pbar.update(len(data)) 

20 return data 

21 

22 def __len__(self): 

23 return self.file_size 

24 

25 def close(self): 

26 self.fp.close() 

27 self.pbar.close() 

28 

29 

30def upload_file_with_retry(bucket_url, file_path, token, max_retries=5): 

31 filename = Path(file_path).name 

32 url = f"{bucket_url}/{filename}" 

33 

34 for attempt in range(max_retries): 

35 try: 

36 print(f"\nAttempt {attempt + 1}/{max_retries}: {filename}") 

37 

38 wrapper = ProgressFileWrapper(file_path) 

39 try: 

40 response = requests.put( 

41 url, 

42 data=wrapper, 

43 headers={'Authorization': f'Bearer {token}'}, 

44 timeout=(30, 300) 

45 ) 

46 response.raise_for_status() 

47 finally: 

48 wrapper.close() 

49 

50 print(f"✓ {filename} uploaded successfully") 

51 return response 

52 

53 except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as e: 

54 print(f"✗ Network error: {e}") 

55 if attempt < max_retries - 1: 

56 wait = 2 ** attempt 

57 print(f"Retrying in {wait}s...") 

58 time.sleep(wait) 

59 else: 

60 raise 

61 except requests.exceptions.HTTPError as e: 

62 print(f"✗ HTTP error: {e}") 

63 raise 

64 

65 

66def main(config_file): 

67 with open(config_file) as f: 

68 config = yaml.safe_load(f) 

69 

70 sandbox = 'sandbox' in config['zenodo_url'] 

71 base_url = 'https://sandbox.zenodo.org/api' if sandbox else 'https://zenodo.org/api' 

72 token = config['access_token'] 

73 project_id = config['project_id'] 

74 

75 response = requests.get( 

76 f"{base_url}/deposit/depositions/{project_id}", 

77 headers={'Authorization': f'Bearer {token}'} 

78 ) 

79 response.raise_for_status() 

80 bucket_url = response.json()['links']['bucket'] 

81 

82 print(f"Project ID: {project_id}") 

83 print(f"Bucket: {bucket_url}") 

84 print(f"Files to upload: {len(config['files'])}") 

85 

86 for file_path in config['files']: 

87 upload_file_with_retry(bucket_url, file_path, token) 

88 

89 

90if __name__ == '__main__': # pragma: no cover 

91 parser = argparse.ArgumentParser() 

92 parser.add_argument('config_file') 

93 args = parser.parse_args() 

94 main(args.config_file)