diff --git a/app/main.py b/app/main.py index 52a3e644..df958326 100644 --- a/app/main.py +++ b/app/main.py @@ -1,5 +1,70 @@ +from typing import List, Optional + + class Cargo: def __init__(self, weight: int) -> None: self.weight = weight -# write your code here + +class BaseRobot: + def __init__(self, + name: str, + weight: int, + coords: Optional[List[int]] = None) -> None: + self.name = name + self.weight = weight + self.coords = coords or [0, 0] + + def go_forward(self, step: int = 1) -> None: + self.coords[1] += step + + def go_back(self, step: int = 1) -> None: + self.coords[1] -= step + + def go_right(self, step: int = 1) -> None: + self.coords[0] += step + + def go_left(self, step: int = 1) -> None: + self.coords[0] -= step + + def get_info(self) -> str: + return f"Robot: {self.name}, Weight: {self.weight}" + + +class FlyingRobot(BaseRobot): + + def __init__(self, + name: str, + weight: int, + coords: Optional[List[int]] = None) -> None: + if coords is None: + coords = [0, 0, 0] + elif len(coords) == 2: + coords.append(0) + super().__init__(name, weight, coords[:2]) + self.coords = coords + + def go_up(self, step: int = 1) -> None: + self.coords[2] += step + + def go_down(self, step: int = 1) -> None: + self.coords[2] -= step + + +class DeliveryDrone(FlyingRobot): + def __init__(self, + name: str, + weight: int, + coords: Optional[List[int]] = None, + max_load_weight: int = 0, + current_load: Cargo = None) -> None: + super().__init__(name, weight, coords) + self.max_load_weight = max_load_weight + self.current_load = current_load + + def hook_load(self, cargo: Cargo) -> None: + if self.current_load is None and cargo.weight <= self.max_load_weight: + self.current_load = cargo + + def unhook_load(self) -> None: + self.current_load = None