From 2ba7edab0c4f866c5ffb7788a4813853878e7899 Mon Sep 17 00:00:00 2001
From: David Lechner <david@pybricks.com>
Date: Mon, 21 Feb 2022 10:55:55 -0600
Subject: [PATCH] pybricks.ev3/pupdevices: fix Motor code completion

The Microsoft Python extension for VS code doesn't include imports in
type hints, so the Motor type wasn't working properly. We can fix this
by wrapping the import type in a new, empty class definition.

The DCMotor class was also missing from ev3devices and the typing for
the overloaded reset_angle method in pupdevices was missing. These are
fixed as well.
---
 CHANGELOG.md                |  5 +++++
 src/pybricks/ev3devices.py  | 12 ++++++++++--
 src/pybricks/ev3devices.pyi |  7 +++++--
 src/pybricks/pupdevices.py  |  8 ++++++--
 src/pybricks/pupdevices.pyi | 18 ++++++++++++++----
 5 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3d19c0d..f30f8fe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,11 @@
 
 ## Unreleased
 
+### Fixed
+- Fixed code completion for `DCMotor` and `Motor` classes in MS Python VS Code extension.
+- Fixed missing `DCMotor` type in `ev3devices`.
+- Fixed type hint for `Motor.reset_angle()` in `pupdevices`.
+
 ## 3.1.0 - 2021-12-16
 
 ### Added
diff --git a/src/pybricks/ev3devices.py b/src/pybricks/ev3devices.py
index be6fdbd..ef266c5 100644
--- a/src/pybricks/ev3devices.py
+++ b/src/pybricks/ev3devices.py
@@ -1,10 +1,18 @@
 # SPDX-License-Identifier: MIT
-# Copyright (c) 2018-2020 The Pybricks Authors
+# Copyright (c) 2018-2022 The Pybricks Authors
 
 """LEGO® MINDSTORMS® EV3 motors and sensors."""
 
 from .parameters import Direction as _Direction
-from ._common import Motor  # noqa E402
+from ._common import DCMotor as _DCMotor, Motor as _Motor
+
+
+class DCMotor(_DCMotor):
+    pass
+
+
+class Motor(_Motor):
+    pass
 
 
 class TouchSensor:
diff --git a/src/pybricks/ev3devices.pyi b/src/pybricks/ev3devices.pyi
index 157fd86..3d27dcc 100644
--- a/src/pybricks/ev3devices.pyi
+++ b/src/pybricks/ev3devices.pyi
@@ -1,10 +1,13 @@
 # SPDX-License-Identifier: MIT
-# Copyright (c) 2020 The Pybricks Authors
+# Copyright (c) 2020-2022 The Pybricks Authors
 
 from typing import List, Optional, Tuple
 
 from .parameters import Color, Direction, Port, Button
-from ._common import Motor  # noqa E402
+from ._common import DCMotor as _DCMotor, Motor as _Motor
+
+class DCMotor(_DCMotor): ...
+class Motor(_Motor): ...
 
 class TouchSensor:
     def __init__(self, port: Port): ...
diff --git a/src/pybricks/pupdevices.py b/src/pybricks/pupdevices.py
index 9ebb396..1a90115 100644
--- a/src/pybricks/pupdevices.py
+++ b/src/pybricks/pupdevices.py
@@ -1,15 +1,19 @@
 # SPDX-License-Identifier: MIT
-# Copyright (c) 2018-2020 The Pybricks Authors
+# Copyright (c) 2018-2022 The Pybricks Authors
 
 """LEGO® Powered Up motor, sensors, and lights."""
 
-from ._common import (Keypad as _Keypad, DCMotor,
+from ._common import (Keypad as _Keypad, DCMotor as _DCMotor,
                       ColorLight as _ColorLight, Motor as _Motor,
                       LightArray as _LightArray, Light as _Light)
 
 from .parameters import Direction as _Direction, Button as _Button
 
 
+class DCMotor(_DCMotor):
+    pass
+
+
 class Motor(_Motor):
     """Generic class to control motors with built-in rotation sensors."""
 
diff --git a/src/pybricks/pupdevices.pyi b/src/pybricks/pupdevices.pyi
index 3776336..97752f8 100644
--- a/src/pybricks/pupdevices.pyi
+++ b/src/pybricks/pupdevices.pyi
@@ -1,13 +1,23 @@
 # SPDX-License-Identifier: MIT
-# Copyright (c) 2020 The Pybricks Authors
+# Copyright (c) 2020-2022 The Pybricks Authors
 
 from typing import Collection, List, Optional, Tuple, Union
 
-from ._common import Keypad, DCMotor, ColorLight, LightArray, Motor as BaseMotor, Light as BaseLight
+from ._common import (
+    Keypad,
+    DCMotor as _DCMotor,
+    ColorLight,
+    LightArray,
+    Motor as _Motor,
+    Light as BaseLight,
+)
 
 from .parameters import Color, Direction, Port
 
-class Motor(BaseMotor): ...
+class DCMotor(_DCMotor): ...
+
+class Motor(_Motor):
+    def reset_angle(self, angle: Optional[int]) -> None: ...
 
 class Remote:
     light: ColorLight
@@ -63,7 +73,7 @@ class ForceSensor:
 class ColorLightMatrix:
     def __init__(self, port: Port) -> None: ...
     def on(self, color: Union[Color, List[Color]]) -> None: ...
-    def off(self)-> None: ...
+    def off(self) -> None: ...
 
 class InfraredSensor:
     def __init__(self, port: Port): ...