atd_db_doc/futool/http/http_request.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2022/8/29 21:31
# @Author  : old tom
# @File    : http_request.py
# @Project : Futool

import urllib.request as req
import urllib.parse
from futool.http.http_response import ResponseWrapper
from http import cookiejar
import json
import shutil
import os

DEFAULT_HEADER = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) "
                  "Chrome/104.0.5112.102 Safari/537.36 Edg/104.0.1293.70",
    "Accept": "*/*",
    "Connection": "keep-alive"
}

DEFAULT_TIMEOUT = 10


class HttpRequestError(Exception):
    def __init__(self, msg=''):
        Exception.__init__(self, msg)


def get(url, param=None, header=None, timeout=DEFAULT_TIMEOUT):
    """
    普通get请求
    :param url: 请求地址
    :param param: 参数字典 自动拼接url
    :param header: 请求头
    :param timeout: 超时(秒)
    :return:
    """
    return request(url, 'GET', param, header=header, timeout=timeout)


def head(url, param=None, header=None, timeout=DEFAULT_TIMEOUT):
    """
    发起head请求
    :param url:
    :param param:
    :param header:
    :param timeout:
    :return:
    """
    return request(url, 'HEAD', param, header, timeout)


def post_form(url, data=None, header=None, timeout=DEFAULT_TIMEOUT):
    """
    post 请求提交普通表单
    :param url:
    :param data:
    :param header:
    :param timeout:
    :return:
    """
    return request(url, 'POST', data=data, header=header, timeout=timeout)


def post(url, json_param=None, header=None, timeout=DEFAULT_TIMEOUT):
    """
    post 请求提交body 参数
    :param url:
    :param json_param:
    :param header:
    :param timeout:
    :return:
    """
    return request(url, 'POST', json_param=json_param, header=header, timeout=timeout)


def request(url, method, param=None, data=None, json_param=None, header=None, timeout=DEFAULT_TIMEOUT, before_send=None,
            success_handler=None,
            error_handler=None):
    """
    基础请求方法
    :param method: 请求方法 目前支持get post 方法
    :param url: 请求地址
    :param param: 请求参数拼接url或表单
    :param data: 表单参数
    :param json_param: json body参数
    :param header: 请求头
    :param timeout: 超时
    :param before_send: 发送前处理
    :param success_handler: 钩子函数（成功处理）
    :param error_handler: 钩子函数（失败处理）
    :return:
    """
    _request = None
    if method not in ['GET', 'POST', 'HEAD']:
        raise HttpRequestError('not support method')
    if header is None:
        header = DEFAULT_HEADER
    opener, cookie = _init_cookie_jar()
    if before_send:
        before_send(opener, cookie, header)
    if method in ['GET', 'HEAD']:
        if param:
            param_str = urllib.parse.urlencode(param)
            url = url + '?' + param_str
        _request = req.Request(url, headers=header, method=method)
    if 'POST' == method:
        if data:
            # 表单
            data = urllib.parse.urlencode(data).encode('utf-8')
            header.update({"Content-Type": "application/x-www-form-urlencoded"})
        if json_param:
            # json body
            data = bytes(json.dumps(json_param), 'utf-8')
            header.update({"Content-Type": "application/json"})
        _request = req.Request(url, data=data, headers=header, method=method)
    wrapper = ResponseWrapper(opener.open(_request, timeout=timeout), cookie)
    if success_handler and wrapper.is_ok():
        return success_handler(wrapper)
    elif error_handler and not wrapper.is_ok():
        return error_handler(wrapper)
    return wrapper


def _init_cookie_jar():
    """
    初始化cookie容器
    :return:
    """
    cookie = cookiejar.CookieJar()
    handler = req.HTTPCookieProcessor(cookie)
    return req.build_opener(handler), cookie


def upload_file():
    # TODO
    pass


def download_file(url, dst_path, overwrite=True, duplicate_handler=None, buffer=16 * 1024):
    """
    文件下载
    :param url: URL路径
    :param dst_path: 目标位置
    :param overwrite: 是否覆盖同名文件
    :param duplicate_handler: 自定义重复文件处理
    :param buffer:
    :return:
    """
    if os.path.isfile(dst_path) and overwrite:
        os.remove(dst_path)
    if os.path.isfile(dst_path) and not overwrite and duplicate_handler:
        duplicate_handler(dst_path)
    b_resp = req.urlopen(url)
    with open(dst_path, 'wb') as f:
        shutil.copyfileobj(b_resp, f, buffer)