Add type-hints to interventions.py
Makes it a bit easier for your editor to perform autocompletion and
goto-definition.
Before this change, it didn't work for me in VS Code (with venv) when trying to
look up `sim.t`, but now it does :tada:
@@ -17,7 +17,7 @@ class Intervention:
self.results = {} #: All interventions are guaranteed to have results, so `Sim` can safely iterate over this dict
- def apply(self, sim) -> None:
+ def apply(self, sim: cv.Sim) -> None:
"""
Apply intervention
@@ -35,7 +35,7 @@ class Intervention:
raise NotImplementedError
- def finalize(self, sim) -> None:
+ def finalize(self, sim: cv.Sim) -> None:
"""
Call function at end of simulation
@@ -50,7 +50,7 @@ class Intervention:
return
- def plot(self, sim, ax) -> None:
+ def plot(self, sim: cv.Sim, ax: pl.Axes) -> None:
"""
Call function during plotting
@@ -120,7 +120,7 @@ class dynamic_pars(Intervention):
return
- def apply(self, sim):
+ def apply(self, sim: cv.Sim):
''' Loop over the parameters, and then loop over the days, applying them if any are found '''
t = sim.t
for parkey,parval in self.pars.items():
@@ -161,13 +161,13 @@ class sequence(Intervention):
return
- def apply(self, sim):
+ def apply(self, sim: cv.Sim):
idx = np.argmax(self._cum_days > sim.t) # Index of the intervention to apply on this day
self.interventions[idx].apply(sim)
return
- def finalize(self, sim, *args, **kwargs):
+ def finalize(self, sim: cv.Sim, *args, **kwargs):
# If any of the sequential interventions write the same quantity to results then
# aggregate them
for intervention in self.interventions:
@@ -208,7 +208,7 @@ class change_beta(Intervention):
return
- def apply(self, sim):
+ def apply(self, sim: cv.Sim):
# If this is the first time it's being run, store beta
if self.orig_beta is None:
@@ -225,7 +225,7 @@ class change_beta(Intervention):
return
- def plot(self, sim, ax):
+ def plot(self, sim: cv.Sim, ax: pl.Axes):
''' Plot vertical lines for when changes in beta '''
ylims = ax.get_ylim()
for day in self.days:
@@ -257,7 +257,7 @@ class test_num(Intervention):
return
- def apply(self, sim):
+ def apply(self, sim: cv.Sim):
t = sim.t
@@ -325,7 +325,7 @@ class test_prob(Intervention):
return
- def apply(self, sim):
+ def apply(self, sim: cv.Sim):
''' Perform testing '''
t = sim.t
@@ -378,7 +378,7 @@ class test_historical(Intervention):
return
- def apply(self, sim):
+ def apply(self, sim: cv.Sim):
''' Perform testing '''
t = sim.t
