improve typing

rebound/horizons.pyi

@@ -3,8 +3,10 @@ from typing import Optional
 from . import Particle
 
 
-def getParticle(particle: Optional[str], m: Optional[float], x: Optional[float], y: Optional[float], z: Optional[float],
-                vx: Optional[float], vy: Optional[float], vz: Optional[float], primary: Optional[float],
-                a: Optional[float], anom: Optional[float], e: Optional[float], omega: Optional[float],
-                inc: Optional[float], Omega: Optional[float], MEAN: Optional[float], date: Optional[float],
-                plane: str, hash: int) -> Particle: ...
+def getParticle(particle: Optional[str] = ..., m: Optional[float] = ..., x: Optional[float] = ...,
+                y: Optional[float] = ..., z: Optional[float] = ..., vx: Optional[float] = ...,
+                vy: Optional[float] = ..., vz: Optional[float] = ..., primary: Optional[float] = ...,
+                a: Optional[float] = ..., anom: Optional[float] = ..., e: Optional[float] = ...,
+                omega: Optional[float] = ..., inc: Optional[float] = ..., Omega: Optional[float] = ...,
+                MEAN: Optional[float] = ..., date: Optional[float] = ..., plane: str = ...,
+                hash: int = ...) -> Particle: ...

rebound/interruptible_pool.pyi

@@ -8,4 +8,4 @@ class InterruptiblePool(Pool):
     def __init__(self, processes: Optional[int] = ..., initializer: Optional[Callable[..., None]] = ...,
                  initargs: Iterable[Any] = ..., **kwargs: Any) -> None: ...
 
-    def map(self, func: Callable, iterable: Iterable, chunksize: Optional[int] = ...): ...
+    def map(self, func: Callable, iterable: Iterable, chunksize: Optional[int] = ...)->Any: ...

rebound/particle.pyi

@@ -29,19 +29,18 @@ class Particle(Structure):
                  primary: Optional[Particle] = ..., a: Optional[float] = ..., P: Optional[float] = ...,
                  e: Optional[float] = ..., inc: Optional[float] = ..., Omega: Optional[float] = ...,
                  omega: Optional[float] = ..., pomega: Optional[float] = ..., f: Optional[float] = ...,
-                 M: Optional[float] = ..., l: Optional[float] = ..., theta: Optional[float] = ...,
-                 T: Optional[float] = ..., r: Optional[float] = ..., date: Optional[str] = ...,
-                 variation: Optional[str] = ..., variation2: Optional[str] = ..., h: Optional[HashType] = ...,
-                 k: Optional[float] = ..., ix: Optional[float] = ..., iy: Optional[float] = ..., hash: HashType = ...,
-                 jacobi_masses: bool = ...): ...
+                 M: Optional[float] = ..., E: Optional[float] = ..., l: Optional[float] = ...,
+                 theta: Optional[float] = ..., T: Optional[float] = ..., r: Optional[float] = ...,
+                 date: Optional[str] = ..., variation: Optional[str] = ..., variation2: Optional[str] = ...,
+                 h: Optional[HashType] = ..., k: Optional[float] = ..., ix: Optional[float] = ...,
+                 iy: Optional[float] = ..., hash: HashType = ..., jacobi_masses: bool = ...) -> None: ...
 
     def copy(self) -> Particle: ...
 
-    def calculate_orbit(self, primary: Optional[Particle], G: Optional[float]) -> Orbit: ...
+    def calculate_orbit(self, primary: Optional[Particle] = ..., G: Optional[float] = ...) -> Orbit: ...
 
-    def sample_orbit(self, Npts: int, primary: Optional[Particle], trailing: bool,
-                     timespan: Optional[float], useTrueAnomaly: Optional[bool],
-                     duplicateEndpoint: Optional[bool]): ... #TODO: return value
+    def sample_orbit(self, Npts: int = ..., primary: Optional[Particle] = ..., samplingAngle: Optional[str] = ...,
+                     duplicateEndpoint: Optional[bool] = ...) -> Any: ...  #TODO: return value
 
     def __pow__(self, other: Particle) -> Particle: ...
 
@@ -170,7 +169,7 @@ class Particle(Structure):
     def orbit(self) -> Orbit: ...
 
     @property
-    def jacobi_com(self): ...  #TODO: correct typing
+    def jacobi_com(self) -> Any: ...  #TODO: correct typing
 
     @property
     def hash(self) -> c_uint32: ...

rebound/plotting.pyi

@@ -20,13 +20,14 @@ def OrbitPlot(sim: Simulation, figsize: Optional[Tuple[float, float]] = ..., fan
 def get_color(color: str) -> ColorTuple: ...
 
 
-def fading_line(x: Iterable[float], y: Iterable[float], color: Union[str, ColorTuple] = ..., alpha_initial: float = ...,
-                alpha_final: float = ..., glow: bool = ..., **kwargs: Any) -> LineCollection: ...
+def fading_line(x: Iterable[float], y: Iterable[float], color: Union[str, ColorTuple] = ..., alpha: float = ...,
+                fading: bool = ..., fancy: bool = ..., **kwargs: Any) -> LineCollection: ...
 
 
 def OrbitPlotOneSlice(sim: Simulation, ax: Axes, Narc: int = ..., color: bool = ..., periastron: bool = ...,
                       orbit_type: str = ..., lw: float = ..., axes: str = ..., plotparticles: List[int] = ...,
-                      primary: Optional[Particle] = ..., fancy: bool = ...) -> None: ...
+                      primary: Optional[Particle] = ..., fancy: bool = ..., xlim: Optional[Tuple[float, float]] = ...,
+                      ylim: Optional[Tuple[float, float]] = ...) -> None: ...
 
 
 def OrbitPlotAddFancyStars(ax: Axes, lw: float, slices: float = ...) -> None: ...

rebound/simulation.pyi

@@ -17,7 +17,7 @@ EOS_TYPES: EnumDict
 BINARY_WARNINGS: List[Tuple[bool, int, str]]
 
 
-class reb_hash__Pointer_pair(Structure):
+class reb_hash_pointer_pair(Structure):
     hash: int
     index: int
 
@@ -169,7 +169,7 @@ class Simulation(Structure):
     var_config: _Pointer[Variation]
     N_active: int
     testparticle_type: int
-    _particle_lookup_table: _Pointer[reb_hash__Pointer_pair]
+    _particle_lookup_table: _Pointer[reb_hash_pointer_pair]
     hash_ctr: int
     N_lookup: int
     allocatedN_lookup: int
@@ -256,9 +256,9 @@ class Simulation(Structure):
     # _extras_cleanup: CFUNCTYPE(None, _Pointer[Simulation])
     extras: c_void_p
 
-    def __new__(cls, *args: Any, **kw: Any)->Simulation: ...
+    def __new__(cls, *args: Any, **kw: Any) -> Simulation: ...
 
-    def __init__(self) -> None: ...
+    def __init__(self, filename: Any = ..., snapshot: Any = ...) -> None: ...
 
     def __repr__(self) -> str: ...
 
@@ -275,7 +275,7 @@ class Simulation(Structure):
 
     visualization: int = ...
 
-    def getWidget(self, **kwargs: Any): ...
+    def getWidget(self, **kwargs: Any) -> Any: ...  #TODO: return type
 
     def refreshWidgets(self) -> None: ...
 
@@ -411,7 +411,7 @@ class Simulation(Structure):
     def convert_particle_units(self, *args: Any) -> None: ...
 
     def add_variation(self, order: int = ..., first_order: Optional[Variation] = ...,
-                      first_order_2: Optional[Variation] = ..., testparticle: int = ...): ...
+                      first_order_2: Optional[Variation] = ..., testparticle: int = ...)->Variation: ...
 
     def init_megno(self, seed: Optional[int] = ...) -> None: ...
 
@@ -435,8 +435,7 @@ class Simulation(Structure):
     def add_particles_ascii(self, s: str) -> None: ...
 
     def calculate_orbits(self, primary: Optional[Particle] = ..., jacobi_masses: bool = ...,
-                         # deprecated options:
-                         # heliocentric: Optional[bool] = ..., barycentric: Optional[bool] = ...
+                         heliocentric: Optional[bool] = ..., barycentric: Optional[bool] = ...
                          ) -> Orbit: ...
 
     def calculate_com(self, first: int = ..., last: Optional[int] = ...) -> Particle: ...

rebound/simulationarchive.pyi

@@ -1,7 +1,7 @@
 from ctypes import Structure, c_int
 from typing import Any, Optional, Callable, Sequence, Iterator, Literal, Iterable, Tuple, Union
 
-import numpy as np
+import numpy
 
 from . import Simulation
 from .types import IntBoolean
@@ -37,4 +37,4 @@ class SimulationArchive(Structure):
 
     def getSimulations(self, times: Iterable[float], **kwargs: Any) -> Iterator[Simulation]: ...
 
-    def getBezierPaths(self, origin: Union[None, int, Literal["com"]] = ...) -> Tuple[np.ndarray, np.ndarray]: ...
+    def getBezierPaths(self, origin: Union[None, int, Literal["com"]] = ...) -> Tuple[numpy.ndarray, numpy.ndarray]: ...