diff --git a/CHANGELOG.md b/CHANGELOG.md index e9905e0e..c99e91c6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### 🎉 New features - Add Postgres execution result store ([#171](https://github.com/roostorg/osprey/pull/171) by [@serendipty01](https://github.com/serendipty01)) + - Add `ParseInt` UDF — converts a numeric string to an integer ([#190](https://github.com/roostorg/osprey/pull/190) by [@bealsbe](https://github.com/bealsbe)) - Add `StringSlice` UDF which extracts a substring by index range ([#189](https://github.com/roostorg/osprey/pull/189) by [@bealsbe](https://github.com/bealsbe)) ### 🐛 Bug fixes diff --git a/osprey_worker/src/osprey/engine/stdlib/udfs/categories.py b/osprey_worker/src/osprey/engine/stdlib/udfs/categories.py index 5b4e40d8..a244cb57 100644 --- a/osprey_worker/src/osprey/engine/stdlib/udfs/categories.py +++ b/osprey_worker/src/osprey/engine/stdlib/udfs/categories.py @@ -13,5 +13,6 @@ class UdfCategories(str, Enum): HTTP = 'HTTP' IP = 'IP' PHONE = 'Phone' + CAST = 'Cast' RANDOM = 'Random' STRING = 'String' diff --git a/osprey_worker/src/osprey/engine/stdlib/udfs/parse_int.py b/osprey_worker/src/osprey/engine/stdlib/udfs/parse_int.py new file mode 100644 index 00000000..e4707606 --- /dev/null +++ b/osprey_worker/src/osprey/engine/stdlib/udfs/parse_int.py @@ -0,0 +1,20 @@ +from osprey.engine.executor.execution_context import ExpectedUdfException + +from ._prelude import ArgumentsBase, ExecutionContext, UDFBase +from .categories import UdfCategories + + +class Arguments(ArgumentsBase): + s: str + + +class ParseInt(UDFBase[Arguments, int]): + """Converts a numeric string to an integer.""" + + category = UdfCategories.CAST + + def execute(self, execution_context: ExecutionContext, arguments: Arguments) -> int: + try: + return int(arguments.s) + except ValueError: + raise ExpectedUdfException() diff --git a/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_parse_int.py b/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_parse_int.py new file mode 100644 index 00000000..299c4e62 --- /dev/null +++ b/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_parse_int.py @@ -0,0 +1,17 @@ +import pytest +from osprey.engine.conftest import ExecuteFunction +from osprey.engine.stdlib.udfs.parse_int import ParseInt +from osprey.engine.udf.registry import UDFRegistry + +pytestmark = [pytest.mark.use_udf_registry(UDFRegistry.with_udfs(ParseInt))] + + +@pytest.mark.parametrize('s,expected', [('15', 15), ('04', 4), ('0', 0), ('-7', -7), ('-092', -92)]) +def test_parse_int(execute: ExecuteFunction, s: str, expected: int) -> None: + data = execute(f'Result = ParseInt(s="{s}")') + assert data == {'Result': expected} + + +def test_parse_int_invalid(execute: ExecuteFunction) -> None: + data = execute('Result = ParseInt(s="ABC")') + assert data == {'Result': None} diff --git a/osprey_worker/src/osprey/worker/_stdlibplugin/udf_register.py b/osprey_worker/src/osprey/worker/_stdlibplugin/udf_register.py index 7bd1e1cb..0325dc1e 100644 --- a/osprey_worker/src/osprey/worker/_stdlibplugin/udf_register.py +++ b/osprey_worker/src/osprey/worker/_stdlibplugin/udf_register.py @@ -20,6 +20,7 @@ from osprey.engine.stdlib.udfs.list_read import ListRead from osprey.engine.stdlib.udfs.list_sort import ListSort from osprey.engine.stdlib.udfs.mx_lookup import MXLookup +from osprey.engine.stdlib.udfs.parse_int import ParseInt from osprey.engine.stdlib.udfs.phone_country import PhoneCountry from osprey.engine.stdlib.udfs.phone_prefix import PhonePrefix from osprey.engine.stdlib.udfs.random_bool import RandomBool @@ -94,6 +95,7 @@ def register_udfs() -> Sequence[Type[UDFBase[Any, Any]]]: ListRead, ListSort, MXLookup, + ParseInt, PhoneCountry, PhonePrefix, RandomBool,