Merge pull request #599 from RocketPy-Team/bug/plot-drag-curves-calla…

…ble-function-source

BUG: plot drag curves when function source is callable

CHANGELOG.md

@@ -53,6 +53,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Fixed
 
+- BUG: plot drag curves when function source is callable [#599](https://github.com/RocketPy-Team/RocketPy/pull/599)
 - BUG: Fix minor type hinting problems [#598](https://github.com/RocketPy-Team/RocketPy/pull/598)
 - BUG: Optional Dependencies Naming in pyproject.toml. [#592](https://github.com/RocketPy-Team/RocketPy/pull/592)
 - BUG: Swap rocket.total_mass.differentiate for motor.total_mass_flow rate [#585](https://github.com/RocketPy-Team/RocketPy/pull/585)

rocketpy/plots/rocket_plots.py

@@ -134,10 +134,22 @@ def drag_curves(self):
         None
         """
 
-        x_power_drag_off = self.rocket.power_off_drag.x_array
-        y_power_drag_off = self.rocket.power_off_drag.y_array
-        x_power_drag_on = self.rocket.power_on_drag.x_array
-        y_power_drag_on = self.rocket.power_on_drag.y_array
+        try:
+            x_power_drag_on = self.rocket.power_on_drag.x_array
+            y_power_drag_on = self.rocket.power_on_drag.y_array
+        except AttributeError:
+            x_power_drag_on = np.linspace(0, 2, 50)
+            y_power_drag_on = np.array(
+                [self.rocket.power_on_drag.source(x) for x in x_power_drag_on]
+            )
+        try:
+            x_power_drag_off = self.rocket.power_off_drag.x_array
+            y_power_drag_off = self.rocket.power_off_drag.y_array
+        except AttributeError:
+            x_power_drag_off = np.linspace(0, 2, 50)
+            y_power_drag_off = np.array(
+                [self.rocket.power_off_drag.source(x) for x in x_power_drag_off]
+            )
 
         fig, ax = plt.subplots()
         ax.plot(x_power_drag_on, y_power_drag_on, label="Power on Drag")