55import shamrock .sys
66
77
8- def purge_old_dumps (dump_prefix , keep_first = 1 , keep_last = 3 , ext = ".sham" ):
8+ def helper_purge_old_dumps (dump_prefix , keep_first = 1 , keep_last = 3 , ext = ".sham" ) -> None :
9+ """
10+ Purge old dump files.
11+ """
912 if shamrock .sys .world_rank () == 0 :
1013 res = glob .glob (dump_prefix + "*" + ext )
1114 res .sort ()
@@ -17,7 +20,10 @@ def purge_old_dumps(dump_prefix, keep_first=1, keep_last=3, ext=".sham"):
1720 os .remove (f )
1821
1922
20- def get_last_dump (dump_prefix , ext = ".sham" ):
23+ def helper_get_last_dump (dump_prefix , ext = ".sham" ) -> int | None :
24+ """
25+ Get the last dump number.
26+ """
2127 res = glob .glob (dump_prefix + "*" + ext )
2228
2329 num_max = - 1
@@ -38,6 +44,10 @@ def get_last_dump(dump_prefix, ext=".sham"):
3844
3945class ShamrockDumpHandleHelper :
4046 def __init__ (self , model , dump_prefix , ext = ".sham" , metadata = False ):
47+ """
48+ Helper class to handle dump files.
49+ """
50+
4151 self .model = model
4252 self .dump_prefix = dump_prefix
4353 self .ext = ext
@@ -47,16 +57,50 @@ def __init__(self, model, dump_prefix, ext=".sham", metadata=False):
4757 def get_dump_name_extension (self , idump , ext ):
4858 return self .dump_prefix + f"{ idump :07} + ext
4959
50- def get_dump_name (self , idump ):
51- return self .get_dump_name_extension (idump , self .ext )
52-
53- def get_last_dump (self ):
54- return get_last_dump (self .dump_prefix , self .ext )
60+ def get_dump_name_extension (self , idump , ext ) -> str :
61+ """Get the name of the dump file with the extension"""
62+ return self .dump_prefix + f"{ idump :07} + ext
5563
56- def purge_old_dumps (self , keep_first = 1 , keep_last = 3 ):
57- purge_old_dumps (self .dump_prefix , keep_first , keep_last , self .ext )
64+ def get_dump_name (self , idump ) -> str :
65+ """Get the name of the dump file (extension from self.ext)"""
66+ return self .get_dump_name_extension (idump , self .ext )
5867
59- def load_dump (self , idump ):
68+ def get_last_dump (self ) -> int | None :
69+ """Find the last dump number"""
70+ return helper_get_last_dump (self .dump_prefix , self .ext )
71+
72+ def purge_old_dumps (self , keep_first = 1 , keep_last = 3 ) -> None :
73+ """
74+ Purge old dump files.
75+
76+ Parameters
77+ ----------
78+ keep_first : int, optional
79+ Number of oldest dump files to keep (default is 1, i.e. keep the first dump).
80+ keep_last : int, optional
81+ Number of newest dump files to keep (default is 3, i.e. keep the last 3 dumps).
82+
83+ Returns
84+ -------
85+ None
86+ This method does not return a value.
87+ """
88+ helper_purge_old_dumps (self .dump_prefix , keep_first , keep_last , self .ext )
89+
90+ def load_dump (self , idump ) -> None :
91+ """
92+ Load a dump file.
93+
94+ Parameters
95+ ----------
96+ idump : int
97+ The dump identifier to load.
98+
99+ Returns
100+ -------
101+ None
102+ This method does not return a value.
103+ """
60104 dump_name = self .get_dump_name (idump )
61105 if shamrock .sys .world_rank () == 0 :
62106 print (f"Loading dump: { dump_name } { idump } )
@@ -69,6 +113,27 @@ def load_dump(self, idump):
69113 return None
70114
71115 def write_dump (self , idump , metadata = None , purge_old_dumps = False , keep_first = 1 , keep_last = 3 ):
116+ """
117+ Write a dump file.
118+
119+ Parameters
120+ ----------
121+ idump : int
122+ The dump identifier to write.
123+ metadata: object
124+ Metadata to be dumped alongside the checkpoint as json
125+ purge_old_dumps : bool, optional
126+ Whether to purge old dumps (default is False).
127+ keep_first : int, optional
128+ Number of oldest dump files to keep (default is 1, i.e. keep the first dump).
129+ keep_last : int, optional
130+ Number of newest dump files to keep (default is 3, i.e. keep the last 3 dumps).
131+
132+ Returns
133+ -------
134+ None
135+ This method does not return a value.
136+ """
72137 dump_name = self .get_dump_name (idump )
73138 self .model .dump (dump_name )
74139
@@ -82,10 +147,19 @@ def write_dump(self, idump, metadata=None, purge_old_dumps=False, keep_first=1,
82147 if purge_old_dumps :
83148 self .purge_old_dumps (keep_first , keep_last )
84149
85- def load_last_dump_or (self , functor_no_last_dump ):
150+ def load_last_dump_or (self , functor_no_last_dump ) -> None :
151+ """
152+ Load the last dump or call a function if no dump is found.
153+
154+ Parameters
155+ ----------
156+ functor_no_last_dump : callable
157+ The function to call if no dump are found (i.e. the setup function).
158+ """
86159 idump = self .get_last_dump ()
87160 if idump is None :
88- functor_no_last_dump ()
89- return None
161+ result = functor_no_last_dump ()
162+ if result is not None :
163+ raise ValueError ("functor_no_last_dump must not return a value" )
90164 else :
91- return self .load_dump (idump )
165+ self .load_dump (idump )
0 commit comments