Lunar (Python): Add type annotations

59_Lunar_LEM_Rocket/python/lunar.py

@@ -6,8 +6,8 @@ Lunar landing simulation
 Ported by Dave LeCompte
 """
 
-import collections
 import math
+from typing import Any, NamedTuple
 
 PAGE_WIDTH = 64
 
@@ -32,7 +32,10 @@ FUEL_RIGHT = FUEL_LEFT + FUEL_WIDTH
 BURN_LEFT = FUEL_RIGHT + COLUMN_WIDTH
 BURN_RIGHT = BURN_LEFT + BURN_WIDTH
 
-PhysicalState = collections.namedtuple("PhysicalState", ["velocity", "altitude"])
+
+class PhysicalState(NamedTuple):
+    velocity: float
+    altitude: float
 
 
 def print_centered(msg: str) -> None:
@@ -42,29 +45,24 @@ def print_centered(msg: str) -> None:
 
 def print_header(title: str) -> None:
     print_centered(title)
-    print_centered("CREATIVE COMPUTING  MORRISTOWN, NEW JERSEY")
-    print()
-    print()
-    print()
+    print_centered("CREATIVE COMPUTING  MORRISTOWN, NEW JERSEY\n\n\n")
 
 
-def add_rjust(line, s, pos):
-    # adds a new field to a line right justified to end at pos
-
-    s = str(s)
-    slen = len(s)
+def add_rjust(line: str, s: Any, pos: int) -> str:
+    """Add a new field to a line right justified to end at pos"""
+    s_str = str(s)
+    slen = len(s_str)
     if len(line) + slen > pos:
         new_len = pos - slen
         line = line[:new_len]
     if len(line) + slen < pos:
         spaces = " " * (pos - slen - len(line))
         line = line + spaces
-    return line + s
+    return line + s_str
 
 
-def add_ljust(line, s, pos):
-    # adds a new field to a line left justified starting at pos
-
+def add_ljust(line: str, s: str, pos: int) -> str:
+    """Add a new field to a line left justified starting at pos"""
     s = str(s)
     if len(line) > pos:
         line = line[:pos]
@@ -75,59 +73,30 @@ def add_ljust(line, s, pos):
 
 
 def print_instructions() -> None:
-    # Somebody had a bad experience with Xerox.
-
+    """Somebody had a bad experience with Xerox."""
     print("THIS IS A COMPUTER SIMULATION OF AN APOLLO LUNAR")
-    print("LANDING CAPSULE.")
-    print()
-    print()
+    print("LANDING CAPSULE.\n\n")
     print("THE ON-BOARD COMPUTER HAS FAILED (IT WAS MADE BY")
-    print("XEROX) SO YOU HAVE TO LAND THE CAPSULE MANUALLY.")
-    print()
+    print("XEROX) SO YOU HAVE TO LAND THE CAPSULE MANUALLY.\n")
 
 
 def print_intro() -> None:
     print("SET BURN RATE OF RETRO ROCKETS TO ANY VALUE BETWEEN")
     print("0 (FREE FALL) AND 200 (MAXIMUM BURN) POUNDS PER SECOND.")
-    print("SET NEW BURN RATE EVERY 10 SECONDS.")
-    print()
-    print("CAPSULE WEIGHT 32,500 LBS; FUEL WEIGHT 16,500 LBS.")
-    print()
-    print()
-    print()
-    print("GOOD LUCK")
-    print()
+    print("SET NEW BURN RATE EVERY 10 SECONDS.\n")
+    print("CAPSULE WEIGHT 32,500 LBS; FUEL WEIGHT 16,500 LBS.\n\n\n")
+    print("GOOD LUCK\n")
 
 
-def show_landing(sim_clock, capsule):
-    w = 3600 * capsule.v
-    print(
-        f"ON MOON AT {sim_clock.elapsed_time:.2f} SECONDS - IMPACT VELOCITY {w:.2f} MPH"
-    )
-    if w < 1.2:
-        print("PERFECT LANDING!")
-    elif w < 10:
-        print("GOOD LANDING (COULD BE BETTER)")
-    elif w <= 60:
-        print("CRAFT DAMAGE... YOU'RE STRANDED HERE UNTIL A RESCUE")
-        print("PARTY ARRIVES. HOPE YOU HAVE ENOUGH OXYGEN!")
-    else:
-        print("SORRY THERE WERE NO SURVIVORS. YOU BLEW IT!")
-        print(f"IN FACT, YOU BLASTED A NEW LUNAR CRATER {w*.227:.2f} FEET DEEP!")
-    end_sim()
-
-
-def show_out_of_fuel(sim_clock, capsule):
-    print(f"FUEL OUT AT {sim_clock.elapsed_time} SECONDS")
-    delta_t = (
-        -capsule.v + math.sqrt(capsule.v**2 + 2 * capsule.a * capsule.g)
-    ) / capsule.g
-    capsule.v += capsule.g * delta_t
-    sim_clock.advance(delta_t)
-    show_landing(sim_clock, capsule)
-
-
-def format_line_for_report(t, miles, feet, velocity, fuel, burn_rate, is_header) -> str:
+def format_line_for_report(
+    t: Any,
+    miles: Any,
+    feet: Any,
+    velocity: Any,
+    fuel: Any,
+    burn_rate: str,
+    is_header: bool,
+) -> str:
     line = add_rjust("", t, SECONDS_RIGHT)
     line = add_rjust(line, miles, ALT_MI_RIGHT)
     line = add_rjust(line, feet, ALT_FT_RIGHT)
@@ -140,42 +109,57 @@ def format_line_for_report(t, miles, feet, velocity, fuel, burn_rate, is_header)
     return line
 
 
+class SimulationClock:
+    def __init__(self, elapsed_time: float, time_until_next_prompt: float) -> None:
+        self.elapsed_time = elapsed_time
+        self.time_until_next_prompt = time_until_next_prompt
+
+    def time_for_prompt(self) -> bool:
+        return self.time_until_next_prompt < 1e-3
+
+    def advance(self, delta_t: float) -> None:
+        self.elapsed_time += delta_t
+        self.time_until_next_prompt -= delta_t
+
+
 class Capsule:
     def __init__(
         self,
-        altitude=120,
-        velocity=1,
-        mass_with_fuel=33000,
-        mass_without_fuel=16500,
-        g=1e-3,
-        z=1.8,
-    ):
+        altitude: float = 120,
+        velocity: float = 1,
+        mass_with_fuel: float = 33000,
+        mass_without_fuel: float = 16500,
+        g: float = 1e-3,
+        z: float = 1.8,
+    ) -> None:
         self.a = altitude  # in miles above the surface
         self.v = velocity  # downward
         self.m = mass_with_fuel
         self.n = mass_without_fuel
         self.g = g
         self.z = z
-        self.fuel_per_second = 0
+        self.fuel_per_second: float = 0
 
-    def remaining_fuel(self):
+    def remaining_fuel(self) -> float:
         return self.m - self.n
 
-    def is_out_of_fuel(self):
+    def is_out_of_fuel(self) -> bool:
         return self.remaining_fuel() < 1e-3
 
-    def update_state(self, sim_clock, delta_t, new_state):
+    def update_state(
+        self, sim_clock: SimulationClock, delta_t: float, new_state: PhysicalState
+    ) -> None:
         sim_clock.advance(delta_t)
         self.m = self.m - delta_t * self.fuel_per_second
         self.a = new_state.altitude
         self.v = new_state.velocity
 
-    def fuel_time_remaining(self):
+    def fuel_time_remaining(self) -> float:
         # extrapolates out how many seconds we have at the current fuel burn rate
         assert self.fuel_per_second > 0
         return self.remaining_fuel() / self.fuel_per_second
 
-    def predict_motion(self, delta_t):
+    def predict_motion(self, delta_t: float) -> PhysicalState:
         # Perform an Euler's Method numerical integration of the equations of motion.
 
         q = delta_t * self.fuel_per_second / self.m
@@ -199,7 +183,7 @@ class Capsule:
 
         return PhysicalState(altitude=new_altitude, velocity=new_velocity)
 
-    def make_state_display_string(self, sim_clock) -> str:
+    def make_state_display_string(self, sim_clock: SimulationClock) -> str:
         seconds = sim_clock.elapsed_time
         miles = int(self.a)
         feet = int(5280 * (self.a - miles))
@@ -211,27 +195,44 @@ class Capsule:
             seconds, miles, feet, velocity, fuel, burn_rate, False
         )
 
-    def prompt_for_burn(self, sim_clock):
+    def prompt_for_burn(self, sim_clock: SimulationClock) -> None:
         msg = self.make_state_display_string(sim_clock)
 
         self.fuel_per_second = float(input(msg))
         sim_clock.time_until_next_prompt = 10
 
 
-class SimulationClock:
-    def __init__(self, elapsed_time, time_until_next_prompt):
-        self.elapsed_time = elapsed_time
-        self.time_until_next_prompt = time_until_next_prompt
-
-    def time_for_prompt(self):
-        return self.time_until_next_prompt < 1e-3
-
-    def advance(self, delta_t):
-        self.elapsed_time += delta_t
-        self.time_until_next_prompt -= delta_t
+def show_landing(sim_clock: SimulationClock, capsule: Capsule) -> None:
+    w = 3600 * capsule.v
+    print(
+        f"ON MOON AT {sim_clock.elapsed_time:.2f} SECONDS - IMPACT VELOCITY {w:.2f} MPH"
+    )
+    if w < 1.2:
+        print("PERFECT LANDING!")
+    elif w < 10:
+        print("GOOD LANDING (COULD BE BETTER)")
+    elif w <= 60:
+        print("CRAFT DAMAGE... YOU'RE STRANDED HERE UNTIL A RESCUE")
+        print("PARTY ARRIVES. HOPE YOU HAVE ENOUGH OXYGEN!")
+    else:
+        print("SORRY THERE WERE NO SURVIVORS. YOU BLEW IT!")
+        print(f"IN FACT, YOU BLASTED A NEW LUNAR CRATER {w*.227:.2f} FEET DEEP!")
+    end_sim()
 
 
-def process_final_tick(delta_t, sim_clock, capsule):
+def show_out_of_fuel(sim_clock: SimulationClock, capsule: Capsule) -> None:
+    print(f"FUEL OUT AT {sim_clock.elapsed_time} SECONDS")
+    delta_t = (
+        -capsule.v + math.sqrt(capsule.v**2 + 2 * capsule.a * capsule.g)
+    ) / capsule.g
+    capsule.v += capsule.g * delta_t
+    sim_clock.advance(delta_t)
+    show_landing(sim_clock, capsule)
+
+
+def process_final_tick(
+    delta_t: float, sim_clock: SimulationClock, capsule: Capsule
+) -> None:
     # When we extrapolated our position based on our velocity
     # and delta_t, we overshot the surface. For better
     # accuracy, we will back up and do shorter time advances.
@@ -255,7 +256,7 @@ def process_final_tick(delta_t, sim_clock, capsule):
         capsule.update_state(sim_clock, delta_t, new_state)
 
 
-def handle_flyaway(sim_clock, capsule):
+def handle_flyaway(sim_clock: SimulationClock, capsule: Capsule) -> bool:
     """
     The user has started flying away from the moon. Since this is a
     lunar LANDING simulation, we wait until the capsule's velocity is
@@ -289,17 +290,11 @@ def handle_flyaway(sim_clock, capsule):
             return False
 
 
-def end_sim():
-    print()
-    print()
-    print()
-    print("TRY AGAIN??")
-    print()
-    print()
-    print()
+def end_sim() -> None:
+    print("\n\n\nTRY AGAIN??\n\n\n")
 
 
-def run_simulation():
+def run_simulation() -> None:
     print()
     print(
         format_line_for_report("SEC", "MI", "FT", "MPH", "LB FUEL", "BURN RATE", True)