Repository: locustio/locust
File: locust/clients.py (Python)

Code before:
        except (MissingSchema, InvalidSchema, InvalidURL):
            raise
        except RequestException as e:
            r = LocustResponse()
            r.error = e
            r.status_code = 0  # with this status_code, content returns None
            r.request = Request(method, url).prepare()
            return r

    # These # type: ignore[override] comments below are needed because our overridden version of functions receives
    # more arguments than functions in the base class.
    def get(self, url: str | bytes, **kwargs: Unpack[RESTKwargs]):  # type: ignore[override]
        kwargs.setdefault("allow_redirects", True)
        return self.request("GET", url, **kwargs)  # type: ignore[misc]

    def options(self, url: str | bytes, **kwargs: Unpack[RESTKwargs]):  # type: ignore[override]
        kwargs.setdefault("allow_redirects", True)
        return self.request("OPTIONS", url, **kwargs)  # type: ignore[misc]

    def head(self, url: str | bytes, **kwargs: Unpack[RESTKwargs]):  # type: ignore[override]
        kwargs.setdefault("allow_redirects", False)
        return self.request("HEAD", url, **kwargs)  # type: ignore[misc]

    # These # type: ignore[misc] comments below are needed because data and json parameters are already defined in the
    # RESTKwargs TypedDict. An alternative approach is to define another TypedDict which doesn't contain them.
    def post(self, url: str | bytes, data: Any | None = None, json: Any | None = None, **kwargs: Unpack[RESTKwargs]):  # type: ignore[override, misc]
        return self.request("POST", url, data=data, json=json, **kwargs)  # type: ignore[misc]

    def put(self, url: str | bytes, data: Any | None = None, **kwargs: Unpack[RESTKwargs]):  # type: ignore[override, misc]
        return self.request("PUT", url, data=data, **kwargs)  # type: ignore[misc]

    def patch(self, url: str | bytes, data: Any | None = None, **kwargs: Unpack[RESTKwargs]):  # type: ignore[override, misc]
        return self.request("PATCH", url, data=data, **kwargs)  # type: ignore[misc]

    def delete(self, url: str | bytes, **kwargs: Unpack[RESTKwargs]):  # type: ignore[override]
        return self.request("DELETE", url, **kwargs)  # type: ignore[misc]


class ResponseContextManager(LocustResponse):
    """
    A Response class that also acts as a context manager that provides the ability to manually
    control if an HTTP request should be marked as successful or a failure in Locust's statistics

    This class is a subclass of :py:class:`Response <requests.Response>` with two additional
    methods: :py:meth:`success <locust.clients.ResponseContextManager.success>` and
    :py:meth:`failure <locust.clients.ResponseContextManager.failure>`.
    """

    _manual_result: bool | Exception | None = None
    _entered = False


Review comment (suggestion):
There does seem to be a problem when using pyright after all. For HTTP methods where  `data=data` and/or `json=json` is passed to `self.request(...)` the return type becomes `Unknown`
![image](https://github.com/locustio/locust/assets/28349560/720c3a67-4d09-497c-a180-f15e7b00a2ea)

I'm unsure what causes this behaviour. Explicitly providing the return type seems to fix it (see suggested changes).

```suggestion
    def post(self, url: str | bytes, data: Any | None = None, json: Any | None = None, **kwargs: Unpack[RESTKwargs]) -> (ResponseContextManager | Response | LocustResponse):  # type: ignore[override, misc]
```

Now the correct return type is shown:
![image](https://github.com/locustio/locust/assets/28349560/25b25044-7a6a-4a56-ac73-fb780c13105a)


Code after:
            return super().request(method, url, **kwargs)
        except (MissingSchema, InvalidSchema, InvalidURL):
            raise
        except RequestException as e:
            r = LocustResponse()
            r.error = e
            r.status_code = 0  # with this status_code, content returns None
            r.request = Request(method, url).prepare()
            return r

    def get(
        self, url: str | bytes, *, data: Any | None = None, json: Any | None = None, **kwargs: Unpack[RESTKwargs]
    ) -> ResponseContextManager | Response | LocustResponse:
        """Sends a GET request"""
        kwargs.setdefault("allow_redirects", True)
        return self.request("GET", url, data=data, json=json, **kwargs)

    def options(
        self,
        url: str | bytes,
        *,
        data: Any | None = None,
        json: Any | None = None,
        **kwargs: Unpack[RESTKwargs],
    ) -> ResponseContextManager | Response | LocustResponse:
        """Sends a OPTIONS request"""
        kwargs.setdefault("allow_redirects", True)
        return self.request("OPTIONS", url, data=data, json=json, **kwargs)

    def head(
        self,
        url: str | bytes,
        *,
        data: Any | None = None,
        json: Any | None = None,
        **kwargs: Unpack[RESTKwargs],
    ) -> ResponseContextManager | Response | LocustResponse:
        """Sends a HEAD request"""
        kwargs.setdefault("allow_redirects", False)
        return self.request("HEAD", url, data=data, json=json, **kwargs)

    def post(
        self,
        url: str | bytes,
        data: Any | None = None,
        json: Any | None = None,
        **kwargs: Unpack[RESTKwargs],
    ) -> ResponseContextManager | Response | LocustResponse:
        """Sends a POST request"""
        return self.request("POST", url, data=data, json=json, **kwargs)