Coverage for /builds/kinetik161/ase/ase/vibrations/pickle2json.py: 81.25%

1import pickle 

2from argparse import ArgumentParser 

3from pathlib import Path 

4 

5import numpy as np 

6 

7from ase.utils.filecache import MultiFileJSONCache 

8 

9description = """ 

10Convert legacy pickle files from ASE vibrations calculations to JSON. 

11 

12Pickles are no longer supported for file storage. 

13 

14WARNING: Only run this command on trusted pickles since unpickling 

15a maliciously crafted pickle allows arbitrary code execution. 

16Indeed that is why pickles are no longer used. 

17""" 

18 

19 

20def port(picklefile): 

21 picklefile = Path(picklefile) 

22 

23 name = picklefile.name 

24 

25 vibname, key, pckl = name.rsplit('.', 3) 

26 assert pckl == 'pckl' 

27 

28 cache = MultiFileJSONCache(picklefile.parent / vibname) 

29 

30 obj = pickle.loads(picklefile.read_bytes()) 

31 if isinstance(obj, np.ndarray): # vibrations 

32 dct = {'forces': obj} 

33 else: # Infrared 

34 forces, dipole = obj 

35 assert isinstance(forces, np.ndarray), f'not supported: {type(forces)}' 

36 assert isinstance(dipole, np.ndarray), f'not supported: {type(dipole)}' 

37 dct = {'forces': forces, 'dipole': dipole} 

38 

39 outfilename = cache._filename(key) 

40 

41 if key in cache: 

42 del cache[key] 

43 

44 cache[key] = dct 

45 print(f'wrote {picklefile} ==> {outfilename}') 

46 

47 

48def main(argv=None): 

49 parser = ArgumentParser(description=description) 

50 parser.add_argument('picklefile', nargs='+') 

51 args = parser.parse_args(argv) 

52 

53 for fname in args.picklefile: 

54 port(fname) 

55 

56 

57if __name__ == '__main__': 

58 main()