python-sdk/robot.py at master · getmaxun/python-sdk · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
from typing import Optional, List
from .client import Client
from .types import MaxunError


class Robot:
    def __init__(self, client: Client, robot_data: dict):
        self.client = client
        self.robot_data = robot_data

    @property
    def id(self) -> str:
        return self.robot_data["recording_meta"]["id"]

    @property
    def name(self) -> str:
        return self.robot_data["recording_meta"]["name"]

    def get_data(self) -> dict:
        return self.robot_data

    async def run(self, options: Optional[dict] = None):
        return await self.client.execute_robot(self.id, options)

    async def get_runs(self) -> list:
        return await self.client.get_runs(self.id)

    async def get_run(self, run_id: str) -> dict:
        return await self.client.get_run(self.id, run_id)

    async def get_run_diff(
        self,
        run_id: str,
        format: Optional[str] = None,
    ) -> dict:
        """Return the detailed monitoring diff for a completed run."""
        return await self.client.get_run_diff(self.id, run_id, format)

    async def get_latest_run(self) -> Optional[dict]:
        runs = await self.get_runs()
        if not runs:
            return None
        return sorted(runs, key=lambda r: r.get("startedAt", ""), reverse=True)[0]

    async def abort(self, run_id: str) -> None:
        await self.client.abort_run(self.id, run_id)

    async def schedule(self, config: dict) -> None:
        updated = await self.client.schedule_robot(self.id, config)
        self.robot_data = updated

    async def unschedule(self) -> None:
        updated = await self.client.unschedule_robot(self.id)
        self.robot_data = updated

    async def add_webhook(self, webhook: dict) -> None:
        updated = await self.client.add_webhook(self.id, webhook)
        self.robot_data = updated

    def get_webhooks(self) -> Optional[list]:
        return self.robot_data.get("webhooks") or None

    async def remove_webhooks(self) -> None:
        updated = await self.client.update_robot(self.id, {"webhooks": None})
        self.robot_data = updated

    def get_schedule(self) -> Optional[dict]:
        return self.robot_data.get("schedule") or None

    async def update(self, updates: dict) -> None:
        updated = await self.client.update_robot(self.id, updates)
        self.robot_data = updated

    async def set_list_limit(self, limit: int) -> None:
        """Set the limit for the first matching scrapeList, crawl, or search action."""
        supported_actions = {"scrapeList", "crawl", "search"}
        workflow = (
            (self.robot_data.get("recording") or {})
            .get("workflow")
            or []
        )

        for pair_index, pair in enumerate(workflow):
            for action_index, action in enumerate(
                pair.get("what") or []
            ):
                if action.get("action") not in supported_actions:
                    continue

                for arg_index, arg in enumerate(
                    action.get("args") or []
                ):
                    if isinstance(arg, dict) and "limit" in arg:
                        self.robot_data = (
                            await self.client.update_list_limits(
                                self.id,
                                [
                                    {
                                        "pairIndex": pair_index,
                                        "actionIndex": action_index,
                                        "argIndex": arg_index,
                                        "limit": limit,
                                    }
                                ],
                            )
                        )
                        return

        raise MaxunError(
            "This robot has no scrapeList, crawl, or search "
            "action with a limit to update."
        )

    async def duplicate(self, target_url: str):
        new_robot_data = await self.client.duplicate_robot(self.id, target_url)
        return Robot(self.client, new_robot_data)

    async def delete(self) -> None:
        await self.client.delete_robot(self.id)

    async def refresh(self) -> None:
        self.robot_data = await self.client.get_robot(self.id)