Skip to content

Commit

Permalink
add SetRPM to python SDK
Browse files Browse the repository at this point in the history
  • Loading branch information
martha-johnston committed May 29, 2024
1 parent 5ae1b82 commit 30a4313
Show file tree
Hide file tree
Showing 6 changed files with 108 additions and 0 deletions.
6 changes: 6 additions & 0 deletions examples/server/v1/components.py
Original file line number Diff line number Diff line change
Expand Up @@ -592,6 +592,12 @@ async def go_to(self, rpm: float, position_revolutions: float, extra: Optional[D
self.position += rps
self.powered = False

async def set_rpm(self, rpm: float, extra: Optional[Dict[str, Any]] = None, **kwargs):
if self.task:
self.task.cancel()
self.powered = True
self.task = asyncio.create_task(self.run_continuously(rpm))

async def reset_zero_position(self, offset: float, extra: Optional[Dict[str, Any]] = None, **kwargs):
if (self.position > 0 and offset > 0) or (self.position < 0 and offset < 0):
self.position = offset
Expand Down
14 changes: 14 additions & 0 deletions src/viam/components/motor/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
MotorServiceStub,
ResetZeroPositionRequest,
SetPowerRequest,
SetRPMRequest,
StopRequest,
)
from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase
Expand Down Expand Up @@ -76,6 +77,19 @@ async def go_to(
request = GoToRequest(name=self.name, rpm=rpm, position_revolutions=position_revolutions, extra=dict_to_struct(extra))
await self.client.GoTo(request, timeout=timeout)

async def set_rpm(
self,
rpm: float,
*,
extra: Optional[Dict[str, Any]] = None,
timeout: Optional[float] = None,
**__,
):
if extra is None:
extra = {}
request = SetRPMRequest(name=self.name, rpm=rpm, extra=dict_to_struct(extra))
await self.client.SetRPM(request, timeout=timeout)

async def reset_zero_position(
self,
offset: float,
Expand Down
25 changes: 25 additions & 0 deletions src/viam/components/motor/motor.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,31 @@ async def go_to(
"""
...

@abc.abstractmethod
async def set_rpm(
self,
rpm: float,
*,
extra: Optional[Dict[str, Any]] = None,
timeout: Optional[float] = None,
**kwargs,
):
"""
Spin the motor indefinitely at the specified speed, in revolutions per minute.
Positive ``rpm`` will result in forward movement and negative ``rpm`` will result in backwards movement
::
my_motor = Motor.from_robot(robot=robot, name="my_motor")
# Spin the motor at 75 RPM.
await my_motor.set_rpm(rpm=75)
Args:
rpm (float): Speed at which the motor should rotate.
"""
...

@abc.abstractmethod
async def reset_zero_position(
self,
Expand Down
11 changes: 11 additions & 0 deletions src/viam/components/motor/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
ResetZeroPositionResponse,
SetPowerRequest,
SetPowerResponse,
SetRPMRequest,
SetRPMResponse,
StopRequest,
StopResponse,
)
Expand Down Expand Up @@ -64,6 +66,15 @@ async def GoTo(self, stream: Stream[GoToRequest, GoToResponse]) -> None:
)
await stream.send_message(GoToResponse())

async def SetRPM(self, stream: Stream[SetRPMRequest, SetRPMResponse]) -> None:
request = await stream.recv_message()
assert request is not None
name = request.name
motor = self.get_resource(name)
timeout = stream.deadline.time_remaining() if stream.deadline else None
await motor.set_rpm(request.rpm, extra=struct_to_dict(request.extra), timeout=timeout, metadata=stream.metadata)
await stream.send_message(SetRPMResponse())

async def ResetZeroPosition(self, stream: Stream[ResetZeroPositionRequest, ResetZeroPositionResponse]) -> None:
request = await stream.recv_message()
assert request is not None
Expand Down
12 changes: 12 additions & 0 deletions tests/mocks/components.py
Original file line number Diff line number Diff line change
Expand Up @@ -701,6 +701,18 @@ async def go_to(
self.extra = extra
self.timeout = timeout

async def set_rpm(
self,
rpm: float,
*,
extra: Optional[Dict[str, Any]] = None,
timeout: Optional[float] = None,
**kwargs,
):
self.powered = rpm != 0
self.extra = extra
self.timeout = timeout

async def reset_zero_position(
self,
offset: float,
Expand Down
40 changes: 40 additions & 0 deletions tests/test_motor.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
MotorServiceStub,
ResetZeroPositionRequest,
SetPowerRequest,
SetRPMRequest,
StopRequest,
)
from viam.resource.manager import ResourceManager
Expand Down Expand Up @@ -83,6 +84,19 @@ async def test_go_to(self, motor: MockMotor):
await motor.go_to(-10, 50)
assert await motor.get_position() == 50

@pytest.mark.asyncio
async def test_set_rpm(self, motor: MockMotor):
await motor.set_rpm(30, timeout=4.56)
assert motor.timeout == loose_approx(4.56)
moving, pwr = await motor.is_powered()
assert moving is True
assert pwr > 0

await motor.set_rpm(-10)
moving, pwr = await motor.is_powered()
assert moving is True
assert pwr < 0

@pytest.mark.asyncio
async def test_reset_zero(self, motor: MockMotor):
await motor.reset_zero_position(20, timeout=5.67)
Expand Down Expand Up @@ -205,6 +219,20 @@ async def test_go_to(self, motor: MockMotor, service: MotorRPCService):
await client.GoTo(request)
assert motor.position == 50

@pytest.mark.asyncio
async def test_set_rpm(self, motor: MockMotor, service: MotorRPCService):
async with ChannelFor([service]) as channel:
client = MotorServiceStub(channel)

request = SetRPMRequest(name=motor.name, rpm=30)
await client.SetRPM(request, timeout=4.56)
assert motor.power > 0
assert motor.timeout == loose_approx(4.56)

request = SetRPMRequest(name=motor.name, rpm=-10)
await client.SetRPM(request)
assert motor.power < 0

@pytest.mark.asyncio
async def test_reset_zero(self, motor: MockMotor, service: MotorRPCService):
async with ChannelFor([service]) as channel:
Expand Down Expand Up @@ -340,6 +368,18 @@ async def test_go_to(self, motor: MockMotor, service: MotorRPCService):
await client.go_to(-10, 50)
assert motor.position == 50

@pytest.mark.asyncio
async def test_set_rpm(self, motor: MockMotor, service: MotorRPCService):
async with ChannelFor([service]) as channel:
client = MotorClient(motor.name, channel)

await client.set_rpm(30, timeout=4.56)
assert motor.power > 0
assert motor.timeout == loose_approx(4.56)

await client.set_rpm(-10)
assert motor.power < 0

@pytest.mark.asyncio
async def test_reset_zero(self, motor: MockMotor, service: MotorRPCService):
async with ChannelFor([service]) as channel:
Expand Down

0 comments on commit 30a4313

Please sign in to comment.