added interface for strains

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## Upcoming
 - Removed the `GameMode` class again; use simple numbers instead (0/1/2/3)
 - Updated to [rosu-pp v0.7](https://github.com/MaxOhn/rosu-pp/blob/main/CHANGELOG.md#v070-2022-07-06)
+- Added `strains` method for `Calculator` which returns instances of `Strains`
 
 # v0.6.0 (2022-07-05)
 - Updated to [PyO3 v0.16](https://github.com/PyO3/pyo3/blob/main/CHANGELOG.md#0165---2022-05-15) from v0.15

README.md

@@ -8,7 +8,7 @@ Check out rosu-pp's README for more info.
 
 ## How to use rosu-pp-py
 
-The library exposes three classes: `Calculator`, `ScoreParams`, and `CalculateResult`. 
+The library exposes three classes: `Calculator`, `ScoreParams`, `CalculateResult`, and `Strains`.
 
 1) The first step is to create a new `Calculator` instance by providing the constructor the path to a `.osu` beatmap file like so
 ```py
@@ -143,6 +143,46 @@ maxCombo: Optional[int]
     The max combo of the map. (O/T/C)
 ```
 
+## Calculating strains
+
+If you want to plot the difficulty of a map over time, you can calculate the strain values.
+The return type of `Calculator::strains` is an instance of the `Strains` class.
+Its attributes depend on the map's game mode again and look as follows:
+```
+sectionLength: float
+    The time in milliseconds between two strain values. (O/T/C/M)
+aim: List[float]
+    Strain values for the aim skill (O)
+aimNoSliders: List[float]
+    Strain values for the aim skill without sliders (O)
+speed: List[float]
+    Strain values for the speed skill (O)
+flashlight: List[float]
+    Strain values for the flashlight skill (O)
+color: List[float]
+    Strain values for the color skill (T)
+rhythm: List[float]
+    Strain values for the rhythm skill (T)
+staminaLeft: List[float]
+    Strain values for the left-stamina skill (T)
+staminaRight: List[float]
+    Strain values for the right-stamina skill (T)
+strains: List[float]
+    Strain values for the strain skill (M)
+movement: List[float]
+    Strain values for the movement skill (C)
+```
+Here's a small example
+```py
+from rosu_pp_py import Calculator
+
+calculator = Calculator('./maps/1980365.osu')
+strains = calculator.strains(8 + 16) # HDHR
+for i,strain in enumerate(strains.aim):
+    currTime = i * strains.sectionLength
+    print(f'Aim strain at {currTime}ms: {strain}')
+```
+
 ## Installing rosu-pp-py
 
 Installing rosu-pp-py requires a [supported version of Python and Rust](https://github.com/PyO3/PyO3#usage).

rosu_pp_py.pyi

@@ -1,4 +1,4 @@
-from typing import Iterable, List, Union
+from typing import Iterable, List, Union, Optional
 
 class ScoreParams:
     """
@@ -48,6 +48,40 @@ class ScoreParams:
     """
     def __init__(self, **kwargs) -> None: ...
 
+class Strains:
+    """
+    A class that contains all strain values of a map.
+    Suitable to plot the difficulty of a map over time.
+    The strain attributes are optional based on the map's mode.
+    In the following, O/T/C/M will denote for which mode the given attribute will be present.
+
+    ## Attributes
+
+    `sectionLength`: float
+        The time in milliseconds between two strain values. (O/T/C/M)
+    `aim`: List[float]
+        Strain values for the aim skill (O)
+    `aimNoSliders`: List[float]
+        Strain values for the aim skill without sliders (O)
+    `speed`: List[float]
+        Strain values for the speed skill (O)
+    `flashlight`: List[float]
+        Strain values for the flashlight skill (O)
+    `color`: List[float]
+        Strain values for the color skill (T)
+    `rhythm`: List[float]
+        Strain values for the rhythm skill (T)
+    `staminaLeft`: List[float]
+        Strain values for the left-stamina skill (T)
+    `staminaRight`: List[float]
+        Strain values for the right-stamina skill (T)
+    `strains`: List[float]
+        Strain values for the strain skill (M)
+    `movement`: List[float]
+        Strain values for the movement skill (C)
+    """
+    def __init__(self) -> None: ...
+
 class CalculateResult:
     """
     A class that contains all difficulty and performance attributes.
@@ -111,7 +145,7 @@ class CalculateResult:
 
 class Calculator:
     """
-    A class to calculate difficulty and performance attributes.
+    A class to calculate difficulty and performance attributes, aswell as strains.
 
     ## Arguments
 
@@ -141,6 +175,8 @@ class Calculator:
         Specify an overall difficulty to override the map's value.
     `calculate(params)`
         Calculate the difficulty and performance attributes for the given score parameters.
+    `strains(mods)`
+        Calculate the strain values for the given mods.
 
     ## Raises
 
@@ -179,4 +215,30 @@ class Calculator:
         # provide params for multiple scores
         results = calculator.calculate([params1, params2])
         ```
+        """
+
+    def strains(self, mods: Optional[int]) -> Strains:
+        """
+        Calculate the strain values for the given mods.
+
+        ## Arguments
+
+        `mods`: Optional[int]
+            Bit value for mods, defaults to 0 (NM) see [https://github.com/ppy/osu-api/wiki#mods](https://github.com/ppy/osu-api/wiki#mods)
+
+        ## Returns
+
+        An instance of the `Strains` class consisting of the strain values
+        for all sections for all skills of the map's game mode,
+        aswell as the section length in milliseconds.
+
+        ## Example
+
+        ```py
+        calculator = Calculator('./maps/1980365.osu')
+        strains = calculator.strains(8 + 16)
+        for i,strain in enumerate(strains.aim):
+            currTime = i * strains.sectionLength
+            print(f'Aim strain at {currTime}ms: {strain}')
+        ```
         """

src/lib.rs

@@ -9,12 +9,13 @@ use pyo3::{
     basic::CompareOp,
     exceptions::{PyException, PyNotImplementedError, PyTypeError},
     prelude::*,
-    types::{PyDict, PyIterator},
+    types::{PyDict, PyIterator, PyTuple},
 };
 use rosu_pp::{
     beatmap::BeatmapAttributes, catch::CatchPerformanceAttributes,
     mania::ManiaPerformanceAttributes, osu::OsuPerformanceAttributes,
     taiko::TaikoPerformanceAttributes, AnyPP, Beatmap, BeatmapExt, GameMode, PerformanceAttributes,
+    Strains as RosuStrains,
 };
 
 #[pyclass]
@@ -87,7 +88,7 @@ impl Calculator {
         self.0.od = od;
     }
 
-    fn calculate(&mut self, py: Python, obj: &PyAny) -> PyResult<Vec<CalculateResult>> {
+    fn calculate(&mut self, obj: &PyAny) -> PyResult<Vec<CalculateResult>> {
         match obj.extract::<ScoreParams>() {
             Ok(params) => {
                 let mods = params.mods;
@@ -102,7 +103,7 @@ impl Calculator {
             Err(_) => {
                 let mut mod_diffs = HashMap::new();
 
-                PyIterator::from_object(py, obj)
+                PyIterator::from_object(obj.py(), obj)
                     .map_err(|_| {
                         let py_type = obj.get_type().name().unwrap_or("<unknown type>");
 
@@ -145,6 +146,105 @@ impl Calculator {
             }
         }
     }
+
+    #[args(args = "*")]
+    fn strains(&mut self, args: &PyTuple) -> PyResult<Strains> {
+        let mods = if let Ok(obj) = args.get_item(0) {
+            if let Ok(mods) = obj.extract::<u32>() {
+                mods
+            } else {
+                let py_type = obj.get_type().name().unwrap_or("<unknown type>");
+                let err = format!("got '{}'; expected 'int'", py_type);
+
+                return Err(PyTypeError::new_err(err));
+            }
+        } else {
+            0
+        };
+
+        Ok(self.0.strains(mods).into())
+    }
+}
+
+#[pyclass]
+#[derive(Clone, Default, PartialEq)]
+#[allow(non_snake_case)]
+struct Strains {
+    #[pyo3(get, set)]
+    sectionLength: f64,
+
+    #[pyo3(get, set)]
+    color: Option<Vec<f64>>,
+    #[pyo3(get, set)]
+    rhythm: Option<Vec<f64>>,
+    #[pyo3(get, set)]
+    staminaLeft: Option<Vec<f64>>,
+    #[pyo3(get, set)]
+    staminaRight: Option<Vec<f64>>,
+
+    #[pyo3(get, set)]
+    aim: Option<Vec<f64>>,
+    #[pyo3(get, set)]
+    aimNoSliders: Option<Vec<f64>>,
+    #[pyo3(get, set)]
+    speed: Option<Vec<f64>>,
+    #[pyo3(get, set)]
+    flashlight: Option<Vec<f64>>,
+
+    #[pyo3(get, set)]
+    strains: Option<Vec<f64>>,
+
+    #[pyo3(get, set)]
+    movement: Option<Vec<f64>>,
+}
+
+#[pymethods]
+impl Strains {
+    fn __richcmp__(&self, other: &PyAny, op: CompareOp) -> PyResult<bool> {
+        match (other.extract::<Self>(), op) {
+            (Ok(ref other), CompareOp::Eq) => Ok(self == other),
+            (Ok(ref other), CompareOp::Ne) => Ok(self != other),
+            _ => Err(PyNotImplementedError::new_err("")),
+        }
+    }
+
+    fn __repr__(&self) -> PyResult<String> {
+        Ok(self.to_string())
+    }
+}
+
+impl From<RosuStrains> for Strains {
+    #[inline]
+    fn from(strains: RosuStrains) -> Self {
+        match strains {
+            RosuStrains::Catch(strains) => Self {
+                sectionLength: strains.section_len,
+                movement: Some(strains.movement),
+                ..Default::default()
+            },
+            RosuStrains::Mania(strains) => Self {
+                sectionLength: strains.section_len,
+                strains: Some(strains.strains),
+                ..Default::default()
+            },
+            RosuStrains::Osu(strains) => Self {
+                sectionLength: strains.section_len,
+                aim: Some(strains.aim),
+                aimNoSliders: Some(strains.aim_no_sliders),
+                speed: Some(strains.speed),
+                flashlight: Some(strains.flashlight),
+                ..Default::default()
+            },
+            RosuStrains::Taiko(strains) => Self {
+                sectionLength: strains.section_len,
+                color: Some(strains.color),
+                rhythm: Some(strains.rhythm),
+                staminaLeft: Some(strains.stamina_left),
+                staminaRight: Some(strains.stamina_right),
+                ..Default::default()
+            },
+        }
+    }
 }
 
 #[pyclass]
@@ -549,6 +649,40 @@ impl ScoreParams {
     }
 }
 
+impl Display for Strains {
+    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
+        let mut s = f.debug_struct("Strains");
+
+        s.field("sectionLength", &self.sectionLength);
+
+        macro_rules! display_field {
+            ($self:ident, $s:ident: $($field:ident,)*) => {
+                $(
+                    if let Some(ref field) = $self.$field {
+                        $s.field(stringify!($field), field);
+                    }
+                )*
+            }
+        }
+
+        display_field! {
+            self, s:
+            color,
+            rhythm,
+            staminaLeft,
+            staminaRight,
+            aim,
+            aimNoSliders,
+            speed,
+            flashlight,
+            strains,
+            movement,
+        }
+
+        s.finish()
+    }
+}
+
 impl Display for CalculateResult {
     fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
         let mut s = f.debug_struct("CalculateResult");
@@ -704,6 +838,7 @@ fn rosu_pp_py(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<ScoreParams>()?;
     m.add_class::<Calculator>()?;
     m.add_class::<CalculateResult>()?;
+    m.add_class::<Strains>()?;
 
     Ok(())
 }