DOC: better docstrings and type hints

Author: Gui-FernandesBR

rocketserializer/cli.py

@@ -191,7 +191,7 @@ def ork2json(
 
         settings = ork_extractor(
             bs=bs,
-            filepath=str(filepath),
+            filepath=filepath,
             output_folder=output,
             ork=ork,
             eng=eng,

rocketserializer/components/drag_curve.py

@@ -1,42 +1,30 @@
 import logging
 import os
 from pathlib import Path
+from typing import Union
 
 import numpy as np
 
 logger = logging.getLogger(__name__)
 
 
-def save_drag_curve(datapoints, data_labels, output_folder):
-    """Extracts the drag curve from the data and saves it to a csv file.
-
-    Parameters
-    ----------
-    datapoints : list
-        The datapoints.
-    data_labels : list of str
-        The labels of the data.
-    path : str
-        Path to the folder where the drag curve should be saved. This needs to
-        be a folder that already exists.
-
-    Returns
-    -------
-    path : str
-        The path to the drag curve.
-    """
-    # Remove the data after apogee
-    altitude_vector = [
+def __collect_altitude_vector(
+    datapoints: list[float], data_labels: list[str]
+) -> list[float]:
+    return [
         float(datapoint.text.split(",")[data_labels.index("Altitude")])
         for datapoint in datapoints
     ]
-    logger.info("Collected altitude vector")
 
+
+def __remove_data_after_apogee(datapoints: list[float], altitude_vector: list[float]):
     apogee_index = np.argmax(altitude_vector)
-    datapoints = datapoints[:apogee_index]
-    logger.info("Removed data after apogee")
+    return datapoints[:apogee_index]
 
-    # Extract the drag coefficient and Mach number
+
+def __extract_drag_and_mach(
+    datapoints: list[float], data_labels: list[str]
+) -> np.ndarray:
     cd = [
         float(datapoint.text.split(",")[data_labels.index("Axial drag coefficient")])
         for datapoint in datapoints
@@ -45,24 +33,41 @@ def save_drag_curve(datapoints, data_labels, output_folder):
         float(datapoint.text.split(",")[data_labels.index("Mach number")])
         for datapoint in datapoints
     ]
-    logger.info("Collected drag coefficient and Mach number")
+    return np.array([mach, cd]).T
+
+
+def __process_drag_data(cd: np.ndarray) -> np.ndarray:
+    cd = cd[~np.isnan(cd).any(axis=1), :]  # Remove NaN values
+    cd = cd[cd[:, 0].argsort()]  # Sort by Mach number
+    cd = np.unique(cd, axis=0)  # Remove duplicate Mach numbers
+    cd = cd[cd[:, 1] > 0, :]  # Remove values when the drag is lower than 0
+    return cd
 
-    # Convert to numpy array
-    cd = np.array([mach, cd]).T
-    # Remove NaN values to avoid errors
-    cd = cd[~np.isnan(cd).any(axis=1), :]
-    # Sort by Mach number
-    cd = cd[cd[:, 0].argsort()]
-    # Remove duplicate Mach numbers
-    cd = np.unique(cd, axis=0)
-    # Remove values when the drag is lower than 0
-    cd = cd[cd[:, 1] > 0, :]
-    logger.info("Successfully created the drag curve")
 
-    # Save to csv file
+def __save_to_csv(cd: np.ndarray, output_folder: Union[Path, str]) -> str:
     path = os.path.join(output_folder, "drag_curve.csv")
     np.savetxt(path, cd, delimiter=",", fmt="%.6f")
     logger.info(
         "Successfully saved the drag curve file to: '%s'", Path(path).as_posix()
     )
     return path
+
+
+def save_drag_curve(
+    datapoints: list[float], data_labels: list[str], output_folder: Union[Path, str]
+) -> str:
+    """Extracts the drag curve from the data and saves it to a csv file."""
+
+    altitude_vector = __collect_altitude_vector(datapoints, data_labels)
+    logger.info("Collected altitude vector")
+
+    datapoints = __remove_data_after_apogee(datapoints, altitude_vector)
+    logger.info("Removed data after apogee")
+
+    cd = __extract_drag_and_mach(datapoints, data_labels)
+    logger.info("Collected drag coefficient and Mach number")
+
+    cd = __process_drag_data(cd)
+    logger.info("Successfully created the drag curve")
+
+    return __save_to_csv(cd, output_folder)

rocketserializer/components/flight.py

@@ -1,19 +1,16 @@
 import logging
 
+from bs4 import BeautifulSoup
+
 from .._helpers import _dict_to_string
 
 logger = logging.getLogger(__name__)
 
 
-def search_launch_conditions(bs):
+def search_launch_conditions(bs: BeautifulSoup) -> dict:
     """Searches the launch conditions in the bs object. Returns a dict with the
     settings.
 
-    Parameters
-    ----------
-    bs : BeautifulSoup
-        The BeautifulSoup object of the open rocket file.
-
     Returns
     -------
     settings : dict

rocketserializer/components/id.py

@@ -1,27 +1,20 @@
 import logging
 from pathlib import Path
 
+from bs4 import BeautifulSoup
+
 from .._helpers import _dict_to_string
 
 logger = logging.getLogger(__name__)
 
 
-def search_id_info(bs, filepath):
+def search_id_info(bs: BeautifulSoup, filepath: Path) -> dict:
     """Searches for the identification of the .ork file
 
-    Parameters
-    ----------
-    bs : BeautifulSoup
-        BeautifulSoup object of the .ork file.
-    filepath : str
-        Path to the .ork file.
-
     Returns
     -------
     dictionary
-        Dictionary with the identification information of the .ork file. The
-        keys are: "rocket_name", "comment", "designer", "ork_version" and
-        "filepath".
+        Dictionary with the identification information of the .ork file.
     """
     settings = {}
     settings["rocket_name"] = bs.find("rocket").find("name").text

rocketserializer/components/motor.py

@@ -1,40 +1,28 @@
 import logging
 import os
 from pathlib import Path
+from typing import Union
 
 import numpy as np
+from bs4 import BeautifulSoup
 
 from .._helpers import _dict_to_string
 
 logger = logging.getLogger(__name__)
 
 
-def search_motor(bs, datapoints, data_labels):
+def search_motor(
+    bs: BeautifulSoup, datapoints: list[float], data_labels: list[str]
+) -> dict:
     """Search for the motor properties in the .ork file. The only property that
     is not included is the thrust curve, which is generated in the
     generate_thrust_curve function. Only rocketpy.SolidMotor class would be able
     to use the information from this function to create a motor object.
 
-    Parameters
-    ----------
-    bs : bs4.BeautifulSoup
-        The BeautifulSoup object of the .ork file.
-    datapoints : list
-        The datapoints from the .ork file.
-    data_labels : list
-        The labels of the datapoints.
-    time_vector : list
-        The time vector of the simulation.
-
     Returns
     -------
     settings : dict
-        Dictionary with the motor properties. The keys are: "burn_time",
-        "grain_density", "grain_initial_inner_radius", "grain_outer_radius",
-        "grain_initial_height", "nozzle_radius", "throat_radius", "dry_mass",
-        "dry_inertia", "center_of_dry_mass", "grains_center_of_mass_position",
-        "grain_number", "grain_separation", "nozzle_position" and
-        "coordinate_system_orientation".
+        Dictionary with the motor properties.
     """
     settings = {}
 
@@ -94,7 +82,12 @@ def search_motor(bs, datapoints, data_labels):
     return settings
 
 
-def generate_thrust_curve(folder_path, datapoints, data_labels, time_vector):
+def generate_thrust_curve(
+    folder_path: Union[str, Path],
+    datapoints: list[float],
+    data_labels: list[str],
+    time_vector: list[float],
+) -> str:
     """Generate the thrust curve from the .ork file.
 
     Parameters
@@ -145,15 +138,10 @@ def generate_thrust_curve(folder_path, datapoints, data_labels, time_vector):
     return source_name
 
 
-def __get_motor_mass(datapoints, data_labels):
-    """Get the motor mass from the .ork file.
-
-    Parameters
-    ----------
-    datapoints : list
-        List of datapoints from the .ork file.
-    data_labels : list
-        List of labels for the datapoints.
+def __get_motor_mass(
+    datapoints: list[float], data_labels: list[str]
+) -> tuple[float, float, int]:
+    """Get the motor mass from the datapoints saved in the .ork file.
 
     Returns
     -------

rocketserializer/components/parachute.py

@@ -1,27 +1,20 @@
 import logging
 
 import numpy as np
+from bs4 import BeautifulSoup
 
 from .._helpers import _dict_to_string
 
 logger = logging.getLogger(__name__)
 
 
-def search_parachutes(bs):
+def search_parachutes(bs: BeautifulSoup) -> dict:
     """Search for the parachutes in the bs and return the settings as a dict.
 
-    Parameters
-    ----------
-    bs : bs4.BeautifulSoup
-        The BeautifulSoup object of the .ork file.
-
     Returns
     -------
     settings : dict
-        A dict containing the settings for the parachutes. The keys are integers
-        and the values are dicts containing the settings for each parachute.
-        The keys of the parachute dicts are: "name", "cd", "cds", "area",
-        "deploy_event", "deploy_delay", "deploy_altitude".
+        A dict containing the settings for the parachutes.
     """
     settings = {}
 

rocketserializer/components/rail_buttons.py

@@ -1,9 +1,11 @@
 import logging
 
+from bs4 import BeautifulSoup
+
 logger = logging.getLogger(__name__)
 
 
-def search_rail_buttons(bs, elements: dict) -> dict:
+def search_rail_buttons(bs: BeautifulSoup, elements: dict) -> dict:
 
     lugs_elements = []
     for value in elements.values():

rocketserializer/components/stored_results.py

@@ -1,25 +1,20 @@
 import logging
 
+from bs4 import BeautifulSoup
+
 from .._helpers import _dict_to_string
 
 logger = logging.getLogger(__name__)
 
 
-def search_stored_results(bs):
+def search_stored_results(bs: BeautifulSoup) -> dict:
     """Search for the stored simulation results in the bs and return the
     settings as a dict.
 
-    Parameters
-    ----------
-    bs : BeautifulSoup
-        BeautifulSoup object of the .ork file.
-
     Returns
     -------
     settings : dict
-        A dict containing the settings for the launch conditions. The keys are:
-        "maxaltitude", "maxvelocity", "maxacceleration", "maxmach", "timetoapogee",
-        "flighttime", "groundhitvelocity" and "launchrodvelocity".
+        A dict containing the settings for the launch conditions.
     """
     settings = {}
 

rocketserializer/components/transition.py

@@ -1,29 +1,23 @@
 import logging
 
+from bs4 import BeautifulSoup
+
 from .._helpers import _dict_to_string
 
 logger = logging.getLogger(__name__)
 
 
-def search_transitions(bs, elements, ork):
-    """Search for the transitions in the bs and return the settings as a dict.
-
-    Parameters
-    ----------
-    bs : bs4.BeautifulSoup
-        The BeautifulSoup object of the .ork file.
-    elements : dict
-        Dictionary with the elements of the rocket.
-    ork : orhelper
-        orhelper object of the open rocket file.
+def search_transitions(
+    bs: BeautifulSoup,
+    elements: dict,
+    ork: "net.sr.openrocket.document.OpenRocketDocument",
+) -> dict:
+    """Search for the transitions in the bs and return the settings as a dict
 
     Returns
     -------
     settings : dict
-        Dictionary with the settings for the transitions. The keys are integers
-        and the values are dicts containing the settings for each transition.
-        The keys of the transition dicts are: "name", "top_radius",
-        "bottom_radius", "length", "position".
+        Dictionary with the settings for the transitions.
     """
     settings = {}
     transitions = bs.findAll("transition")

rocketserializer/nb_builder.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+from pathlib import Path
 
 import nbformat as nbf
 
@@ -12,7 +13,7 @@ class NotebookBuilder:
     using rocketpy simulation on it
     """
 
-    def __init__(self, parameters_json: str) -> None:
+    def __init__(self, parameters_json: Path) -> None:
         """read the file and process the dictionary do not build anything yet"""
         self.parameters_json = parameters_json
         self.trapezoidal_fins_check = False