MyPy errors, refactor

src/poliastro/plotting/_base.py

@@ -1,5 +1,7 @@
 import os.path
+from collections import namedtuple
 from itertools import cycle
+from typing import List
 
 import numpy as np
 import plotly.colors
@@ -12,6 +14,10 @@
 from poliastro.util import norm
 
 
+class Trajectory(namedtuple("Trajectory", ["trajectory", "state", "label", "color"])):
+    pass
+
+
 class _BaseOrbitPlotter:
     """
     Parent Class for the 2D and 3D OrbitPlotter Classes based on Plotly.
@@ -21,7 +27,7 @@ def __init__(self, figure=None):
         self._figure = figure or FigureWidget()
         self._layout = None
 
-        self._trajectories = []
+        self._trajectories = []  # type: List[Trajectory]
 
         self._attractor = None
         self._attractor_radius = np.inf * u.km

src/poliastro/plotting/misc.py

@@ -1,3 +1,5 @@
+from typing import Union
+
 from poliastro.bodies import (
     Earth,
     Jupiter,
@@ -22,23 +24,24 @@ def plot_solar_system(outer=True, epoch=None, use_3d=False):
     ------------
     outer : bool, optional
         Whether to print the outer Solar System, default to True.
-    epoch: ~astropy.time.Time, optional
+    epoch : ~astropy.time.Time, optional
         Epoch value of the plot, default to J2000.
+    use_3d : bool, optional
+        Produce 3D plot, default to False.
+
     """
     bodies = [Mercury, Venus, Earth, Mars]
     if outer:
         bodies.extend([Jupiter, Saturn, Uranus, Neptune])
 
     if use_3d:
-        op = OrbitPlotter3D()
+        op = OrbitPlotter3D()  # type: Union[OrbitPlotter3D, OrbitPlotter2D]
     else:
         op = OrbitPlotter2D()
+        op.set_frame(*Orbit.from_body_ephem(Earth, epoch).pqw())  # type: ignore
 
     for body in bodies:
         orb = Orbit.from_body_ephem(body, epoch)
         op.plot(orb, label=str(body))
 
-    if not use_3d:
-        op.set_frame(*Orbit.from_body_ephem(Earth, epoch).pqw())
-
     return op

src/poliastro/plotting/static.py

@@ -1,4 +1,3 @@
-from collections import namedtuple
 from typing import List
 
 import matplotlib as mpl
@@ -11,9 +10,7 @@
 from poliastro.twobody.propagation import mean_motion
 from poliastro.util import norm
 
-
-class _Trajectory(namedtuple("_Trajectory", ["trajectory", "state", "label", "color"])):
-    pass
+from ._base import Trajectory
 
 
 class StaticOrbitPlotter(object):
@@ -50,7 +47,7 @@ def __init__(self, ax=None, num_points=150, dark=False):
         self._frame = None
         self._attractor = None
         self._attractor_radius = np.inf * u.km
-        self._trajectories = []  # type: List[_Trajectory]
+        self._trajectories = []  # type: List[Trajectory]
 
     @property
     def trajectories(self):
@@ -113,7 +110,7 @@ def plot_trajectory(self, trajectory, *, label=None, color=None):
             )
 
         self._trajectories.append(
-            _Trajectory(trajectory, None, label, lines[0].get_color())
+            Trajectory(trajectory, None, label, lines[0].get_color())
         )
 
         return lines
@@ -205,7 +202,7 @@ def plot(self, orbit, label=None, color=None, method=mean_motion):
         lines = self._plot(positions, orbit.r, label, color)
 
         self._trajectories.append(
-            _Trajectory(positions, orbit.r, label, lines[0].get_color())
+            Trajectory(positions, orbit.r, label, lines[0].get_color())
         )
 
         return lines