webscout 1.2.1__py3-none-any.whl → 1.2.2__py3-none-any.whl

DeepWEBS/utilsdw/logger.py

@@ -0,0 +1,269 @@
+import datetime
+import functools
+import inspect
+import logging
+import os
+import shutil
+import subprocess
+from termcolor import colored
+
+
+def add_fillers(text, filler="=", fill_side="both"):
+    terminal_width = shutil.get_terminal_size().columns
+    text = text.strip()
+    text_width = len(text)
+    if text_width >= terminal_width:
+        return text
+
+    if fill_side[0].lower() == "b":
+        leading_fill_str = filler * ((terminal_width - text_width) // 2 - 1) + " "
+        trailing_fill_str = " " + filler * (
+            terminal_width - text_width - len(leading_fill_str) - 1
+        )
+    elif fill_side[0].lower() == "l":
+        leading_fill_str = filler * (terminal_width - text_width - 1) + " "
+        trailing_fill_str = ""
+    elif fill_side[0].lower() == "r":
+        leading_fill_str = ""
+        trailing_fill_str = " " + filler * (terminal_width - text_width - 1)
+    else:
+        raise ValueError("Invalid fill_side")
+
+    filled_str = f"{leading_fill_str}{text}{trailing_fill_str}"
+    return filled_str
+
+
+class OSLogger(logging.Logger):
+    LOG_METHODS = {
+        "err": ("error", "red"),
+        "warn": ("warning", "light_red"),
+        "note": ("info", "light_magenta"),
+        "mesg": ("info", "light_cyan"),
+        "file": ("info", "light_blue"),
+        "line": ("info", "white"),
+        "success": ("info", "light_green"),
+        "fail": ("info", "light_red"),
+        "back": ("debug", "light_cyan"),
+    }
+    INDENT_METHODS = [
+        "indent",
+        "set_indent",
+        "reset_indent",
+        "store_indent",
+        "restore_indent",
+        "log_indent",
+    ]
+    LEVEL_METHODS = [
+        "set_level",
+        "store_level",
+        "restore_level",
+        "quiet",
+        "enter_quiet",
+        "exit_quiet",
+    ]
+    LEVEL_NAMES = {
+        "critical": logging.CRITICAL,
+        "error": logging.ERROR,
+        "warning": logging.WARNING,
+        "info": logging.INFO,
+        "debug": logging.DEBUG,
+    }
+
+    def __init__(self, name=None, prefix=False):
+        if not name:
+            frame = inspect.stack()[1]
+            module = inspect.getmodule(frame[0])
+            name = module.__name__
+
+        super().__init__(name)
+        self.setLevel(logging.INFO)
+
+        if prefix:
+            formatter_prefix = "[%(asctime)s] - [%(name)s] - [%(levelname)s]\n"
+        else:
+            formatter_prefix = ""
+
+        self.formatter = logging.Formatter(formatter_prefix + "%(message)s")
+
+        stream_handler = logging.StreamHandler()
+        stream_handler.setLevel(logging.INFO)
+        stream_handler.setFormatter(self.formatter)
+        self.addHandler(stream_handler)
+
+        self.log_indent = 0
+        self.log_indents = []
+
+        self.log_level = "info"
+        self.log_levels = []
+
+    def indent(self, indent=2):
+        self.log_indent += indent
+
+    def set_indent(self, indent=2):
+        self.log_indent = indent
+
+    def reset_indent(self):
+        self.log_indent = 0
+
+    def store_indent(self):
+        self.log_indents.append(self.log_indent)
+
+    def restore_indent(self):
+        self.log_indent = self.log_indents.pop(-1)
+
+    def set_level(self, level):
+        self.log_level = level
+        self.setLevel(self.LEVEL_NAMES[level])
+
+    def store_level(self):
+        self.log_levels.append(self.log_level)
+
+    def restore_level(self):
+        self.log_level = self.log_levels.pop(-1)
+        self.set_level(self.log_level)
+
+    def quiet(self):
+        self.set_level("critical")
+
+    def enter_quiet(self, quiet=False):
+        if quiet:
+            self.store_level()
+            self.quiet()
+
+    def exit_quiet(self, quiet=False):
+        if quiet:
+            self.restore_level()
+
+    def log(
+        self,
+        level,
+        color,
+        msg,
+        indent=0,
+        fill=False,
+        fill_side="both",
+        end="\n",
+        *args,
+        **kwargs,
+    ):
+        if type(msg) == str:
+            msg_str = msg
+        else:
+            msg_str = repr(msg)
+            quotes = ["'", '"']
+            if msg_str[0] in quotes and msg_str[-1] in quotes:
+                msg_str = msg_str[1:-1]
+
+        indent_str = " " * (self.log_indent + indent)
+        indented_msg = "\n".join([indent_str + line for line in msg_str.split("\n")])
+
+        if fill:
+            indented_msg = add_fillers(indented_msg, fill_side=fill_side)
+
+        handler = self.handlers[0]
+        handler.terminator = end
+
+        getattr(self, level)(colored(indented_msg, color), *args, **kwargs)
+
+    def route_log(self, method, msg, *args, **kwargs):
+        level, method = method
+        functools.partial(self.log, level, method, msg)(*args, **kwargs)
+
+    def err(self, msg: str = "", *args, **kwargs):
+        self.route_log(("error", "red"), msg, *args, **kwargs)
+
+    def warn(self, msg: str = "", *args, **kwargs):
+        self.route_log(("warning", "light_red"), msg, *args, **kwargs)
+
+    def note(self, msg: str = "", *args, **kwargs):
+        self.route_log(("info", "light_magenta"), msg, *args, **kwargs)
+
+    def mesg(self, msg: str = "", *args, **kwargs):
+        self.route_log(("info", "light_cyan"), msg, *args, **kwargs)
+
+    def file(self, msg: str = "", *args, **kwargs):
+        self.route_log(("info", "light_blue"), msg, *args, **kwargs)
+
+    def line(self, msg: str = "", *args, **kwargs):
+        self.route_log(("info", "white"), msg, *args, **kwargs)
+
+    def success(self, msg: str = "", *args, **kwargs):
+        self.route_log(("info", "light_green"), msg, *args, **kwargs)
+
+    def fail(self, msg: str = "", *args, **kwargs):
+        self.route_log(("info", "light_red"), msg, *args, **kwargs)
+
+    def back(self, msg: str = "", *args, **kwargs):
+        self.route_log(("debug", "light_cyan"), msg, *args, **kwargs)
+
+
+logger = OSLogger()
+
+
+def shell_cmd(cmd, getoutput=False, showcmd=True, env=None):
+    if showcmd:
+        logger.info(colored(f"\n$ [{os.getcwd()}]", "light_blue"))
+        logger.info(colored(f"  $ {cmd}\n", "light_cyan"))
+    if getoutput:
+        output = subprocess.getoutput(cmd, env=env)
+        return output
+    else:
+        subprocess.run(cmd, shell=True, env=env)
+
+
+class Runtimer:
+    def __enter__(self):
+        self.t1, _ = self.start_time()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.t2, _ = self.end_time()
+        self.elapsed_time(self.t2 - self.t1)
+
+    def start_time(self):
+        t1 = datetime.datetime.now()
+        self.logger_time("start", t1)
+        return t1, self.time2str(t1)
+
+    def end_time(self):
+        t2 = datetime.datetime.now()
+        self.logger_time("end", t2)
+        return t2, self.time2str(t2)
+
+    def elapsed_time(self, dt=None):
+        if dt is None:
+            dt = self.t2 - self.t1
+        self.logger_time("elapsed", dt)
+        return dt, self.time2str(dt)
+
+    def logger_time(self, time_type, t):
+        time_types = {
+            "start": "Start",
+            "end": "End",
+            "elapsed": "Elapsed",
+        }
+        time_str = add_fillers(
+            colored(
+                f"{time_types[time_type]} time: [ {self.time2str(t)} ]",
+                "light_magenta",
+            ),
+            fill_side="both",
+        )
+        logger.line(time_str)
+
+    # Convert time to string
+    def time2str(self, t):
+        datetime_str_format = "%Y-%m-%d %H:%M:%S"
+        if isinstance(t, datetime.datetime):
+            return t.strftime(datetime_str_format)
+        elif isinstance(t, datetime.timedelta):
+            hours = t.seconds // 3600
+            hour_str = f"{hours} hr" if hours > 0 else ""
+            minutes = (t.seconds // 60) % 60
+            minute_str = f"{minutes:>2} min" if minutes > 0 else ""
+            seconds = t.seconds % 60
+            second_str = f"{seconds:>2} s"
+            time_str = " ".join([hour_str, minute_str, second_str]).strip()
+            return time_str
+        else:
+            return str(t)

webscout/DWEBS.py

@@ -0,0 +1,179 @@
+
+from pydantic import BaseModel, Field
+from typing import Union
+
+from DeepWEBS.utilsdw.logger import logger
+from DeepWEBS.networks.google_searcher import GoogleSearcher
+from DeepWEBS.networks.webpage_fetcher import BatchWebpageFetcher
+from DeepWEBS.documents.query_results_extractor import QueryResultsExtractor
+from DeepWEBS.documents.webpage_content_extractor import BatchWebpageContentExtractor
+from DeepWEBS.utilsdw.logger import logger
+import argparse
+
+class DeepWEBS:
+    def __init__(self):
+        pass
+
+    class DeepSearch(BaseModel):
+        queries: list = Field(
+            default=[""],
+            description="(list[str]) Queries to search",
+        )
+        result_num: int = Field(
+            default=10,
+            description="(int) Number of search results",
+        )
+        safe: bool = Field(
+            default=False,
+            description="(bool) Enable SafeSearch",
+        )
+        types: list = Field(
+            default=["web"],
+            description="(list[str]) Types of search results: `web`, `image`, `videos`, `news`",
+        )
+        extract_webpage: bool = Field(
+            default=False,
+            description="(bool) Enable extracting main text contents from webpage, will add `text` filed in each `query_result` dict",
+        )
+        overwrite_query_html: bool = Field(
+            default=False,
+            description="(bool) Overwrite HTML file of query results",
+        )
+        overwrite_webpage_html: bool = Field(
+            default=False,
+            description="(bool) Overwrite HTML files of webpages from query results",
+        )
+
+    def queries_to_search_results(self, item: DeepSearch):
+        google_searcher = GoogleSearcher()
+        queries_search_results = []
+        for query in item.queries:
+            query_results_extractor = QueryResultsExtractor()
+            if not query.strip():
+                continue
+            query_html_path = google_searcher.search(
+                query=query,
+                result_num=item.result_num,
+                safe=item.safe,
+                overwrite=item.overwrite_query_html,
+            )
+            query_search_results = query_results_extractor.extract(query_html_path)
+            queries_search_results.append(query_search_results)
+        logger.note(queries_search_results)
+
+        if item.extract_webpage:
+            queries_search_results = self.extract_webpages(
+                queries_search_results,
+                overwrite_webpage_html=item.overwrite_webpage_html,
+            )
+        return queries_search_results
+
+    def extract_webpages(self, queries_search_results, overwrite_webpage_html=False):
+        for query_idx, query_search_results in enumerate(queries_search_results):
+            # Fetch webpages with urls
+            batch_webpage_fetcher = BatchWebpageFetcher()
+            urls = [
+                query_result["url"]
+                for query_result in query_search_results["query_results"]
+            ]
+            url_and_html_path_list = batch_webpage_fetcher.fetch(
+                urls,
+                overwrite=overwrite_webpage_html,
+                output_parent=query_search_results["query"],
+            )
+
+            # Extract webpage contents from htmls
+            html_paths = [
+                str(url_and_html_path["html_path"])
+                for url_and_html_path in url_and_html_path_list
+            ]
+            batch_webpage_content_extractor = BatchWebpageContentExtractor()
+            html_path_and_extracted_content_list = (
+                batch_webpage_content_extractor.extract(html_paths)
+            )
+
+            # Build the map of url to extracted_content
+            html_path_to_url_dict = {
+                str(url_and_html_path["html_path"]): url_and_html_path["url"]
+                for url_and_html_path in url_and_html_path_list
+            }
+            url_to_extracted_content_dict = {
+                html_path_to_url_dict[
+                    html_path_and_extracted_content["html_path"]
+                ]: html_path_and_extracted_content["extracted_content"]
+                for html_path_and_extracted_content in html_path_and_extracted_content_list
+            }
+
+            # Write extracted contents (as 'text' field) to query_search_results
+            for query_result_idx, query_result in enumerate(
+                query_search_results["query_results"]
+            ):
+                url = query_result["url"]
+                extracted_content = url_to_extracted_content_dict[url]
+                queries_search_results[query_idx]["query_results"][query_result_idx][
+                    "text"
+                ] = extracted_content
+
+        return queries_search_results
+
+
+class ArgParser(argparse.ArgumentParser):
+    def __init__(self, *args, **kwargs):
+        super(ArgParser, self).__init__(*args, **kwargs)
+
+        self.add_argument(
+            "-q",
+            "--queries",
+            type=str,
+            nargs="+",
+            required=True,
+            help="Queries to search",
+        )
+        self.add_argument(
+            "-n",
+            "--result_num",
+            type=int,
+            default=10,
+            help="Number of search results",
+        )
+        self.add_argument(
+            "-s",
+            "--safe",
+            default=False,
+            action="store_true",
+            help="Enable SafeSearch",
+        )
+        self.add_argument(
+            "-t",
+            "--types",
+            type=str,
+            nargs="+",
+            default=["web"],
+            choices=["web", "image", "videos", "news"],
+            help="Types of search results",
+        )
+        self.add_argument(
+            "-e",
+            "--extract_webpage",
+            default=False,
+            action="store_true",
+            help="Enable extracting main text contents from webpage",
+        )
+        self.add_argument(
+            "-o",
+            "--overwrite_query_html",
+            default=False,
+            action="store_true",
+            help="Overwrite HTML file of query results",
+        )
+        self.add_argument(
+            "-w",
+            "--overwrite_webpage_html",
+            default=False,
+            action="store_true",
+            help="Overwrite HTML files of webpages from query results",
+        )
+
+        self.args = self.parse_args()
+
+

webscout/__init__.py

@@ -8,7 +8,7 @@ import logging
 from .webscout_search import WEBS
 from .webscout_search_async import AsyncWEBS
 from .version import __version__
-
+from .DWEBS import DeepWEBS
 __all__ = ["WEBS", "AsyncWEBS", "__version__", "cli"]
 
 logging.getLogger("webscout").addHandler(logging.NullHandler())

webscout/version.py

@@ -1,2 +1,2 @@
-__version__ = "1.2.1"
+__version__ = "1.2.2"