feat: add build method back to multihosturl (#730)

ollz272 · davidhewitt · web-flow · commit 8adeccff0cde · 2023-07-11T10:35:57.000+01:00
Co-authored-by: Oli Parker &lt;oliver.parker@gridedge.co.uk&gt;
Co-authored-by: David Hewitt &lt;1939362+davidhewitt@users.noreply.github.com&gt;
diff --git a/python/pydantic_core/_pydantic_core.pyi b/python/pydantic_core/_pydantic_core.pyi
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import decimal
 import sys
-from typing import Any, Callable, Generic, Type, TypeVar
+from typing import Any, Callable, Generic, Optional, Type, TypeVar
 
 from pydantic_core import ErrorDetails, ErrorTypeInfo, InitErrorDetails, MultiHostHost
 from pydantic_core.core_schema import CoreConfig, CoreSchema, ErrorType
@@ -191,6 +191,19 @@ class Url(SupportsAllComparisons):
     def __repr__(self) -> str: ...
     def __str__(self) -> str: ...
     def __deepcopy__(self, memo: dict) -> str: ...
+    @classmethod
+    def build(
+        cls,
+        *,
+        scheme: str,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        host: str,
+        port: Optional[int] = None,
+        path: Optional[str] = None,
+        query: Optional[str] = None,
+        fragment: Optional[str] = None,
+    ) -> str: ...
 
 class MultiHostUrl(SupportsAllComparisons):
     def __new__(cls, url: str) -> Self: ...
@@ -208,6 +221,20 @@ class MultiHostUrl(SupportsAllComparisons):
     def __repr__(self) -> str: ...
     def __str__(self) -> str: ...
     def __deepcopy__(self, memo: dict) -> Self: ...
+    @classmethod
+    def build(
+        cls,
+        *,
+        scheme: str,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        host: Optional[str] = None,
+        hosts: Optional[list[MultiHostHost]] = None,
+        port: Optional[int] = None,
+        path: Optional[str] = None,
+        query: Optional[str] = None,
+        fragment: Optional[str] = None,
+    ) -> str: ...
 
 @final
 class SchemaError(Exception):
diff --git a/src/url.rs b/src/url.rs
@@ -1,13 +1,17 @@
 use std::collections::hash_map::DefaultHasher;
+use std::fmt;
+use std::fmt::Formatter;
 use std::hash::{Hash, Hasher};
 
 use idna::punycode::decode_to_string;
+use pyo3::exceptions::PyValueError;
 use pyo3::once_cell::GILOnceCell;
-use pyo3::prelude::*;
 use pyo3::pyclass::CompareOp;
-use pyo3::types::PyDict;
+use pyo3::types::{PyDict, PyType};
+use pyo3::{intern, prelude::*};
 use url::Url;
 
+use crate::tools::SchemaDict;
 use crate::SchemaValidator;
 
 static SCHEMA_DEFINITION_URL: GILOnceCell<SchemaValidator> = GILOnceCell::new();
@@ -150,6 +154,42 @@ impl PyUrl {
     fn __getnewargs__(&self) -> (&str,) {
         (self.__str__(),)
     }
+
+    #[classmethod]
+    #[pyo3(signature=(*, scheme, host, username=None, password=None, port=None, path=None, query=None, fragment=None))]
+    #[allow(clippy::too_many_arguments)]
+    pub fn build<'a>(
+        cls: &'a PyType,
+        scheme: &str,
+        host: &str,
+        username: Option<&str>,
+        password: Option<&str>,
+        port: Option<u16>,
+        path: Option<&str>,
+        query: Option<&str>,
+        fragment: Option<&str>,
+    ) -> PyResult<&'a PyAny> {
+        let url_host = UrlHostParts {
+            username: username.map(Into::into),
+            password: password.map(Into::into),
+            host: Some(host.into()),
+            port,
+        };
+        let mut url = format!("{scheme}://{url_host}");
+        if let Some(path) = path {
+            url.push('/');
+            url.push_str(path);
+        }
+        if let Some(query) = query {
+            url.push('?');
+            url.push_str(query);
+        }
+        if let Some(fragment) = fragment {
+            url.push('#');
+            url.push_str(fragment);
+        }
+        cls.call1((url,))
+    }
 }
 
 #[pyclass(name = "MultiHostUrl", module = "pydantic_core._pydantic_core", subclass)]
@@ -314,6 +354,114 @@ impl PyMultiHostUrl {
     fn __getnewargs__(&self) -> (String,) {
         (self.__str__(),)
     }
+
+    #[classmethod]
+    #[pyo3(signature=(*, scheme, hosts=None, path=None, query=None, fragment=None, host=None, username=None, password=None, port=None))]
+    #[allow(clippy::too_many_arguments)]
+    pub fn build<'a>(
+        cls: &'a PyType,
+        scheme: &str,
+        hosts: Option<Vec<UrlHostParts>>,
+        path: Option<&str>,
+        query: Option<&str>,
+        fragment: Option<&str>,
+        // convenience parameters to build with a single host
+        host: Option<&str>,
+        username: Option<&str>,
+        password: Option<&str>,
+        port: Option<u16>,
+    ) -> PyResult<&'a PyAny> {
+        let mut url =
+            if hosts.is_some() && (host.is_some() || username.is_some() || password.is_some() || port.is_some()) {
+                return Err(PyValueError::new_err(
+                    "expected one of `hosts` or singular values to be set.",
+                ));
+            } else if let Some(hosts) = hosts {
+                // check all of host / user / password / port empty
+                // build multi-host url
+                let mut multi_url = format!("{scheme}://");
+                for (index, single_host) in hosts.iter().enumerate() {
+                    if single_host.is_empty() {
+                        return Err(PyValueError::new_err(
+                            "expected one of 'host', 'username', 'password' or 'port' to be set",
+                        ));
+                    }
+                    multi_url.push_str(&single_host.to_string());
+                    if index != hosts.len() - 1 {
+                        multi_url.push(',');
+                    };
+                }
+                multi_url
+            } else if host.is_some() {
+                let url_host = UrlHostParts {
+                    username: username.map(Into::into),
+                    password: password.map(Into::into),
+                    host: host.map(Into::into),
+                    port: port.map(Into::into),
+                };
+                format!("{scheme}://{url_host}")
+            } else {
+                return Err(PyValueError::new_err("expected either `host` or `hosts` to be set"));
+            };
+
+        if let Some(path) = path {
+            url.push('/');
+            url.push_str(path);
+        }
+        if let Some(query) = query {
+            url.push('?');
+            url.push_str(query);
+        }
+        if let Some(fragment) = fragment {
+            url.push('#');
+            url.push_str(fragment);
+        }
+        cls.call1((url,))
+    }
+}
+
+pub struct UrlHostParts {
+    username: Option<String>,
+    password: Option<String>,
+    host: Option<String>,
+    port: Option<u16>,
+}
+
+impl UrlHostParts {
+    fn is_empty(&self) -> bool {
+        self.host.is_none() && self.password.is_none() && self.host.is_none() && self.port.is_none()
+    }
+}
+
+impl FromPyObject<'_> for UrlHostParts {
+    fn extract(ob: &'_ PyAny) -> PyResult<Self> {
+        let py = ob.py();
+        let dict = ob.downcast::<PyDict>()?;
+        Ok(UrlHostParts {
+            username: dict.get_as(intern!(py, "username"))?,
+            password: dict.get_as(intern!(py, "password"))?,
+            host: dict.get_as(intern!(py, "host"))?,
+            port: dict.get_as(intern!(py, "port"))?,
+        })
+    }
+}
+
+impl fmt::Display for UrlHostParts {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        match (&self.username, &self.password) {
+            (Some(username), None) => write!(f, "{username}@")?,
+            (None, Some(password)) => write!(f, ":{password}@")?,
+            (Some(username), Some(password)) => write!(f, "{username}:{password}@")?,
+            (None, None) => {}
+        };
+        if let Some(host) = &self.host {
+            write!(f, "{host}")?;
+        }
+        if let Some(port) = self.port {
+            write!(f, ":{port}")?;
+        }
+        Ok(())
+    }
 }
 
 fn host_to_dict<'a>(py: Python<'a>, lib_url: &Url) -> PyResult<&'a PyDict> {
diff --git a/tests/validators/test_url.py b/tests/validators/test_url.py
@@ -1216,3 +1216,86 @@ def test_url_hash() -> None:
 
 def test_url_deepcopy() -> None:
     assert deepcopy(Url('http://example.com')) == Url('http://example.com/')
+
+
+def test_multi_url_build() -> None:
+    url = MultiHostUrl.build(
+        scheme='postgresql',
+        username='testuser',
+        password='testpassword',
+        host='127.0.0.1',
+        port=5432,
+        path='database',
+        query='sslmode=require',
+        fragment='test',
+    )
+    assert url == MultiHostUrl('postgresql://testuser:testpassword@127.0.0.1:5432/database?sslmode=require#test')
+    assert str(url) == 'postgresql://testuser:testpassword@127.0.0.1:5432/database?sslmode=require#test'
+
+
+@pytest.mark.parametrize('field', ['host', 'password', 'username', 'port'])
+def test_multi_url_build_hosts_set_with_single_value(field) -> None:
+    """Hosts can't be provided with any single url values."""
+    hosts = [
+        {'host': '127.0.0.1', 'password': 'testpassword', 'username': 'testuser', 'port': 5432},
+        {'host': '127.0.0.1', 'password': 'testpassword', 'username': 'testuser', 'port': 5432},
+    ]
+    kwargs = dict(scheme='postgresql', hosts=hosts, path='database', query='sslmode=require', fragment='test')
+    if field == 'port':
+        kwargs[field] = 5432
+    else:
+        kwargs[field] = 'test'
+    with pytest.raises(ValueError):
+        MultiHostUrl.build(**kwargs)
+
+
+def test_multi_url_build_hosts_empty_host() -> None:
+    """Hosts can't be provided with any single url values."""
+    hosts = [{}]
+    with pytest.raises(ValueError):
+        MultiHostUrl.build(scheme='postgresql', hosts=hosts, path='database', query='sslmode=require', fragment='test')
+
+
+def test_multi_url_build_hosts() -> None:
+    """Hosts can't be provided with any single url values."""
+    hosts = [
+        {'host': '127.0.0.1', 'password': 'testpassword', 'username': 'testuser', 'port': 5431},
+        {'host': '127.0.0.1', 'password': 'testpassword', 'username': 'testuser', 'port': 5433},
+    ]
+    kwargs = dict(scheme='postgresql', hosts=hosts, path='database', query='sslmode=require', fragment='test')
+    url = MultiHostUrl.build(**kwargs)
+    assert url == MultiHostUrl(
+        'postgresql://testuser:testpassword@127.0.0.1:5431,testuser:testpassword@127.0.0.1:5433/database?sslmode=require#test'
+    )
+    assert (
+        str(url)
+        == 'postgresql://testuser:testpassword@127.0.0.1:5431,testuser:testpassword@127.0.0.1:5433/database?sslmode=require#test'
+    )
+
+
+def test_multi_url_build_neither_host_and_hosts_set() -> None:
+    with pytest.raises(ValueError):
+        MultiHostUrl.build(
+            scheme='postgresql',
+            username='testuser',
+            password='testpassword',
+            port=5432,
+            path='database',
+            query='sslmode=require',
+            fragment='test',
+        )
+
+
+def test_url_build() -> None:
+    url = Url.build(
+        scheme='postgresql',
+        username='testuser',
+        password='testpassword',
+        host='127.0.0.1',
+        port=5432,
+        path='database',
+        query='sslmode=require',
+        fragment='test',
+    )
+    assert url == Url('postgresql://testuser:testpassword@127.0.0.1:5432/database?sslmode=require#test')
+    assert str(url) == 'postgresql://testuser:testpassword@127.0.0.1:5432/database?sslmode=require#test'
