explicitly list input particle types in docstrings

Author: joeljonsson

pyapr/tree/fill_tree.py

@@ -27,15 +27,15 @@ def fill_tree_mean(apr: APR,
     ----------
     apr: APR
         Input APR data structure.
-    parts: ParticleData
+    parts: ByteParticles, ShortParticles, FloatParticles or LongParticles
         Input APR particle values.
-    output: ParticleData, optional
-        (optional) Particle data object for the output values. If provided, the type must either be FloatParticles
+    output: ByteParticles, ShortParticles, FloatParticles or LongParticles, optional
+        Particle data object for the output values. If provided, the type must either be FloatParticles
         or the same type as ``parts``. If None, generates a new object of the same type as ``parts``. (default: None)
 
     Returns
     -------
-    output: ParticleData
+    output: ByteParticles, ShortParticles, FloatParticles or LongParticles
         The computed tree values.
     """
     _check_input(apr, parts)
@@ -54,15 +54,15 @@ def fill_tree_max(apr: APR,
     ----------
     apr: APR
         Input APR data structure.
-    parts: ParticleData
+    parts: ByteParticles, ShortParticles, FloatParticles or LongParticles
         Input APR particle values.
-    output: ParticleData, optional
-        (optional) Particle data object for the output values. If provided, the type must either be FloatParticles
+    output: ByteParticles, ShortParticles, FloatParticles or LongParticles, optional
+        Particle data object for the output values. If provided, the type must either be FloatParticles
         or the same type as ``parts``. If None, generates a new object of the same type as ``parts``. (default: None)
 
     Returns
     -------
-    output: ParticleData
+    output: ByteParticles, ShortParticles, FloatParticles or LongParticles
         The computed tree values.
     """
     _check_input(apr, parts)
@@ -81,15 +81,15 @@ def fill_tree_min(apr: APR,
     ----------
     apr: APR
         Input APR data structure.
-    parts: ParticleData
+    parts: ByteParticles, ShortParticles, FloatParticles or LongParticles
         Input APR particle values.
-    output: ParticleData, optional
-        (optional) Particle data object for the output values. If provided, the type must either be FloatParticles
+    output: ByteParticles, ShortParticles, FloatParticles or LongParticles, optional
+        Particle data object for the output values. If provided, the type must either be FloatParticles
         or the same type as ``parts``. If None, generates a new object of the same type as ``parts``. (default: None)
 
     Returns
     -------
-    output: ParticleData
+    output: ByteParticles, ShortParticles, FloatParticles or LongParticles
         The computed tree values.
     """
     _check_input(apr, parts)

pyapr/tree/tree_ops.py

@@ -19,9 +19,9 @@ def sample_from_tree(apr: APR,
     ----------
     apr: APR
         Input APR data structure.
-    parts: ByteParticles, ShortParticles, FloatParticles, LongParticles
+    parts: ByteParticles, ShortParticles, FloatParticles or LongParticles
         Input APR particle values.
-    tree_parts: ByteParticles, ShortParticles, FloatParticles, LongParticles
+    tree_parts: ByteParticles, ShortParticles, FloatParticles or LongParticles
         Input APR tree values. Must either be of type FloatParticles or the same type as `parts`.
     num_levels: int
         Sample values from level `apr.level_max()-num_levels`. If `num_levels=1`, each particle takes the value
@@ -32,7 +32,7 @@ def sample_from_tree(apr: APR,
 
     Returns
     -------
-    output: ByteParticles, ShortParticles, FloatParticles, LongParticles
+    output: ByteParticles, ShortParticles, FloatParticles or LongParticles
         The resampled particle values.
     """
     _check_input(apr, parts)