Add payload mass

2 files changed, 6 insertions, 3 deletions

diff --git a/main.py b/main.py
index d2c92d1..93a1304 100644
--- a/main.py
+++ b/main.py
@@ -40,7 +40,8 @@ def main(argv):
                         )
 
     rocket = Rocket(name="starship launch system", 
-                    stages=[first_stage, second_stage]
+                    stages=[first_stage, second_stage],
+                    payload_mass=100
                     )
     
     body = Body(name="earth",
@@ -96,6 +97,7 @@ def main(argv):
         #TODO: do max load on rocket so it blows up
         #TODO: allow for x movement, speed, accel etc
         #TODO: allow multilanguage api for landing algorithms etc
+        #TODO: probs need cloud sprite that spawns and despawns as well as floor sprite
 
 def linear_gradient(start_color, end_color, length, value_at):
     return [
diff --git a/rocket.py b/rocket.py
index 97aaa59..962678b 100644
--- a/rocket.py
+++ b/rocket.py
@@ -1,10 +1,11 @@
 from stage import Stage
 
 class Rocket():
-    def __init__(self, name: str, stages: [type[Stage]]):
+    def __init__(self, name: str, stages: [type[Stage]], payload_mass: int):
         self.name = name
         self.stages = stages
         self.stages_spent = 0
+        self.payload_mass = payload_mass
 
     def current_stage(self) -> type[Stage]:
         return self.stages[0]
@@ -19,7 +20,7 @@ class Rocket():
             self.current_stage().engines_on = engines_on
 
     def total_mass(self):
-        total_mass = 0
+        total_mass = self.payload_mass
         for stage in self.stages:
             total_mass += stage.total_mass()
         return total_mass