diff --git a/app/main.py b/app/main.py index 52a3e644..3dfb304f 100644 --- a/app/main.py +++ b/app/main.py @@ -2,4 +2,67 @@ class Cargo: def __init__(self, weight: int) -> None: self.weight = weight -# write your code here + +class BaseRobot: + def __init__( + self, + name: str, + weight: int | float, + coords: list = None + ) -> None: + self.coords = [0, 0] if coords is None else coords + self.name = name + self.weight = weight + + 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 | float, + coords: list = None + ) -> None: + super().__init__(name, weight, coords) + self.coords = [0, 0, 0] if coords is None else 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 | float, + max_load_weight: int | float, + current_load: None | Cargo, + coords: list = None + ) -> None: + super().__init__(name, weight, coords) + self.current_load = current_load + self.max_load_weight = max_load_weight + + def hook_load(self, cargo: Cargo) -> None: + if (not self.current_load) and (cargo.weight <= self.max_load_weight): + self.current_load = cargo + + def unhook_load(self) -> None: + self.current_load = None