grab.spider.queue_backend.memory — Grab 0.6.41 documentation

# Source code for grab.spider.queue_backend.memory

    from __future__ import annotations

    from contextlib import suppress
    from datetime import datetime
    from queue import Empty, PriorityQueue

    from grab.spider.queue_backend.base import BaseTaskQueue
    from grab.spider.task import Task

    class MemoryTaskQueue(BaseTaskQueue):
        def __init__(self) -> None:
            super().__init__()
            self.queue_object: PriorityQueue[tuple[int, Task]] = PriorityQueue()
            self.schedule_list: list[tuple[datetime, Task]] = []

        def put(
            self, task: Task, priority: int, schedule_time: None | datetime = None
        ) -> None:
            if schedule_time is None:
                self.queue_object.put((priority, task))
            else:
                self.schedule_list.append((schedule_time, task))

        def get(self) -> Task:
            now = datetime.utcnow()

            removed_indexes = []
            for idx, item in enumerate(self.schedule_list):
                schedule_time, task = item
                if schedule_time <= now:
                    self.put(task, 1)
                    removed_indexes.append(idx)

            self.schedule_list = [
                x for idx, x in enumerate(self.schedule_list) if idx not in removed_indexes
            ]

            _, task = self.queue_object.get(block=False)
            return task

        def size(self) -> int:
            return self.queue_object.qsize() + len(self.schedule_list)

        def clear(self) -> None:
            with suppress(Empty):
                while True:
                    self.queue_object.get(False)
            self.schedule_list = []

        def close(self) -> None:
            pass