apify · jirimoravcik · Jan 4, 2023 · Dec 19, 2022 · Dec 19, 2022 · Dec 19, 2022
diff --git a/mypy.ini b/mypy.ini
@@ -1,4 +1,5 @@
 [mypy]
+python_version=3.8
 files =
    docs,
    scripts,

diff --git a/pytest.ini b/pytest.ini
@@ -0,0 +1,2 @@
+[pytest]
+asyncio_mode=auto
diff --git a/setup.py b/setup.py
@@ -57,6 +57,8 @@
        'pyee ~= 9.0.4',
        'typing-extensions ~= 4.4.0',
        'websockets ~= 10.4',
+        'aiofiles ~= 22.1.0',
+        'aioshutil ~= 1.2',
    ],
    extras_require={
        'dev': [
@@ -79,6 +81,7 @@
            'sphinx ~= 5.3.0',
            'sphinx-autodoc-typehints ~= 1.19.5',
            'sphinx-markdown-builder == 0.5.4',  # pinned to 0.5.4, because 0.5.5 has a formatting bug
+            'types-aiofiles ~= 22.1.0.4',
            'types-psutil ~= 5.9.5.5',
            'types-setuptools ~= 65.6.0.1',
        ],

diff --git a/src/apify/_types.py b/src/apify/_types.py
@@ -0,0 +1,6 @@
+from typing import Any, Dict, List, Union
+
+# Type for representing json-serializable values
+# It's close enough to the real thing supported by json.parse, and the best we can do until mypy supports recursive types
+# It was suggested in a discussion with (and approved by) Guido van Rossum, so I'd consider it correct enough
+JSONSerializable = Union[str, int, float, bool, None, Dict[str, Any], List[Any]]
diff --git a/src/apify/_utils.py b/src/apify/_utils.py
@@ -1,13 +1,23 @@
 import asyncio
+import base64
+import contextlib
+import hashlib
 import inspect
+import io
+import json
+import mimetypes
 import os
+import re
 import sys
 import time
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Any, Callable, Generic, Optional, TypeVar, Union, overload
+from typing import Any, Callable, Dict, Generic, List, NoReturn, Optional, TypeVar, Union, cast, overload

+import aioshutil
 import psutil
+from aiofiles import ospath
+from aiofiles.os import remove, rename

 from apify_client import __version__ as client_version

@@ -20,7 +30,9 @@
    BOOL_ENV_VARS,
    DATETIME_ENV_VARS,
    INTEGER_ENV_VARS,
+    REQUEST_ID_LENGTH,
    ApifyEnvVars,
+    StorageTypes,
 )


@@ -164,3 +176,137 @@ async def _run_func_at_interval_async(func: Callable, interval_secs: float) -> N
        res = func()
        if inspect.isawaitable(res):
            await res
+
+
+class ListPage:  # TODO: Rather use exported version from Apify client
+    """A single page of items returned from a list() method."""
+
+    #: list: List of returned objects on this page
+    items: List
+    #: int: Count of the returned objects on this page
+    count: int
+    #: int: The limit on the number of returned objects offset specified in the API call
+    offset: int
+    #: int: The offset of the first object specified in the API call
+    limit: int
+    #: int: Total number of objects matching the API call criteria
+    total: int
+    #: bool: Whether the listing is descending or not
+    desc: bool
+
+    def __init__(self, data: Dict) -> None:
+        """Initialize a ListPage instance from the API response data."""
+        self.items = data['items'] if 'items' in data else []
+        self.offset = data['offset'] if 'offset' in data else 0
+        self.limit = data['limit'] if 'limit' in data else 0
+        self.count = data['count'] if 'count' in data else len(self.items)
+        self.total = data['total'] if 'total' in data else self.offset + self.count
+        self.desc = data['desc'] if 'desc' in data else False
+
+
+async def _force_remove(filename: str) -> None:
+    """JS-like rm(filename, { force: true })."""
+    with contextlib.suppress(FileNotFoundError):
+        await remove(filename)
+
+
+def _json_serializer(obj: Any) -> str:  # TODO: Decide how to parse/dump/handle datetimes!
+    if isinstance(obj, (datetime)):
+        return obj.isoformat(timespec='milliseconds') + 'Z'
+    else:
+        return str(obj)
+
+
+def _filter_out_none_values_recursively(dictionary: Dict) -> Dict:
+    """Return copy of the dictionary, recursively omitting all keys for which values are None."""
+    return cast(dict, _filter_out_none_values_recursively_internal(dictionary))
+
+
+# Unfortunately, it's necessary to have an internal function for the correct result typing, without having to create complicated overloads
+def _filter_out_none_values_recursively_internal(dictionary: Dict, remove_empty_dicts: Optional[bool] = None) -> Optional[Dict]:
+    result = {}
+    for k, v in dictionary.items():
+        if isinstance(v, dict):
+            v = _filter_out_none_values_recursively_internal(v, remove_empty_dicts is True or remove_empty_dicts is None)
+        if v is not None:
+            result[k] = v
+    if not result and remove_empty_dicts:
+        return None
+    return result
+
+
+def _json_dumps(obj: Any) -> str:
+    """Dump JSON to a string with the correct settings and serializer."""
+    return json.dumps(obj, ensure_ascii=False, indent=2, default=_json_serializer)
+
+
+uuid_regex = re.compile('[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}', re.I)
+
+
+def _is_uuid(string: str) -> bool:
+    """Test whether the given string matches UUID format."""
+    return bool(uuid_regex.match(string))
+
+
+def _raise_on_non_existing_storage(client_type: StorageTypes, id: str) -> NoReturn:
+    raise ValueError(f'{client_type} with id: {id} does not exist.')
+
+
+def _raise_on_duplicate_storage(client_type: StorageTypes, key_name: str, value: str) -> NoReturn:
+    raise ValueError(f'{client_type} with {key_name}: {value} already exists.')
+
+
+def _guess_file_extension(content_type: str) -> Optional[str]:
+    """Guess the file extension based on content type."""
+    # e.g. mimetypes.guess_extension('application/json ') does not work...
+    actual_content_type = content_type.split(';')[0].strip()
+    ext = mimetypes.guess_extension(actual_content_type)
+    # Remove the leading dot if extension successfully parsed
+    return ext[1:] if ext is not None else ext
+
+
+def _is_content_type_json(content_type: str) -> bool:
+    return bool(re.search(r'^application/json', content_type, flags=re.IGNORECASE))
+
+
+def _is_content_type_xml(content_type: str) -> bool:
+    return bool(re.search(r'^application/.*xml$', content_type, flags=re.IGNORECASE))
+
+
+def _is_content_type_text(content_type: str) -> bool:
+    return bool(re.search(r'^text/', content_type, flags=re.IGNORECASE))
+
+
+def _is_file_or_bytes(value: Any) -> bool:
+    # The check for IOBase is not ideal, it would be better to use duck typing,
+    # but then the check would be super complex, judging from how the 'requests' library does it.
+    # This way should be good enough for the vast majority of use cases, if it causes issues, we can improve it later.
+    return isinstance(value, (bytes, bytearray, io.IOBase))
+
+
+def _maybe_parse_body(body: bytes, content_type: str) -> Any:  # TODO: Improve return type
+    try:
+        if _is_content_type_json(content_type):
+            return json.loads(body)  # Returns any
+        elif _is_content_type_xml(content_type) or _is_content_type_text(content_type):
+            return body.decode('utf-8')  # TODO: Check if utf-8 can be assumed
+    except ValueError as err:
+        print('_maybe_parse_body error', err)
+    return body
+
+
+def _unique_key_to_request_id(unique_key: str) -> str:
+    """Generate request ID based on unique key in a deterministic way."""
+    id = re.sub(r'(\+|\/|=)', '', base64.b64encode(hashlib.sha256(unique_key.encode('utf-8')).digest()).decode('utf-8'))
+
+    return id[:REQUEST_ID_LENGTH] if len(id) > REQUEST_ID_LENGTH else id
+
+
+async def _force_rename(src_dir: str, dst_dir: str) -> None:
+    """Rename a directory. Checks for existence of soruce directory and removes destination directory if it exists."""
+    # Make sure source directory exists
+    if await ospath.exists(src_dir):
+        # Remove destination directory if it exists
+        if await ospath.exists(dst_dir):
+            await aioshutil.rmtree(dst_dir, ignore_errors=True)
+        await rename(src_dir, dst_dir)
diff --git a/src/apify/actor.py b/src/apify/actor.py
@@ -1,7 +1,7 @@
 import asyncio
-import datetime
 import functools
 import inspect
+from datetime import datetime
 from types import TracebackType
 from typing import Any, Awaitable, Callable, Coroutine, Dict, List, Optional, Type, TypeVar, Union, cast

@@ -216,7 +216,7 @@ def _get_system_info(self) -> Dict:
        memory_usage_bytes = _get_memory_usage_bytes()
        # This is in camel case to be compatible with the events from the platform
        result = {
-            'createdAt': datetime.datetime.now().isoformat(timespec='milliseconds') + 'Z',
+            'createdAt': datetime.utcnow().isoformat(timespec='milliseconds') + 'Z',
            'cpuCurrentUsage': cpu_usage_percent,
            'memCurrentBytes': memory_usage_bytes,
        }

diff --git a/src/apify/consts.py b/src/apify/consts.py
@@ -133,3 +133,16 @@ class ApifyEnvVars(str, Enum):
 ]

 STRING_ENV_VARS: List[_STRING_ENV_VARS_TYPE] = list(get_args(_STRING_ENV_VARS_TYPE))
+
+
+class StorageTypes(str, Enum):
+    """Possible Apify storage types."""
+
+    DATASET = 'Dataset'
+    KEY_VALUE_STORE = 'Key-value store'
+    REQUEST_QUEUE = 'Request queue'
+
+
+DEFAULT_API_PARAM_LIMIT = 1000
+
+REQUEST_ID_LENGTH = 15
diff --git a/src/apify/memory_storage/file_storage_utils.py b/src/apify/memory_storage/file_storage_utils.py
@@ -0,0 +1,108 @@
+import os
+from typing import Dict, List, Tuple
+
+import aiofiles
+from aiofiles.os import makedirs
+
+from .._utils import _force_remove, _json_dumps
+
+
+async def _update_metadata(*, data: Dict, entity_directory: str, write_metadata: bool) -> None:
+    # Skip writing the actual metadata file. This is done after ensuring the directory exists so we have the directory present
+    if not write_metadata:
+        return
+
+    # Ensure the directory for the entity exists
+    await makedirs(entity_directory, exist_ok=True)
+
+    # Write the metadata to the file
+    file_path = os.path.join(entity_directory, '__metadata__.json')
+    async with aiofiles.open(file_path, mode='wb') as f:
+        await f.write(_json_dumps(data).encode('utf-8'))
+
+
+async def _update_dataset_items(
+    *,
+    data: List[Tuple[str, Dict]],
+    entity_directory: str,
+    persist_storage: bool,
+) -> None:
+    # Skip writing files to the disk if the client has the option set to false
+    if not persist_storage:
+        return
+
+    # Ensure the directory for the entity exists
+    await makedirs(entity_directory, exist_ok=True)
+
+    # Save all the new items to the disk
+    for idx, item in data:
+        file_path = os.path.join(entity_directory, f'{idx}.json')
+        async with aiofiles.open(file_path, mode='wb') as f:
+            await f.write(_json_dumps(item).encode('utf-8'))
+
+
+async def _set_or_delete_key_value_store_record(
+    *,
+    entity_directory: str,
+    persist_storage: bool,
+    record: Dict,
+    should_set: bool,
+    write_metadata: bool,
+) -> None:
+    # Skip writing files to the disk if the client has the option set to false
+    if not persist_storage:
+        return
+
+    # Ensure the directory for the entity exists
+    await makedirs(entity_directory, exist_ok=True)
+
+    # Create files for the record
+    record_path = os.path.join(entity_directory, f"""{record['key']}.{record['extension']}""")
+    record_metadata_path = os.path.join(entity_directory, f"""{record['key']}.__metadata__.json""")
+
+    await _force_remove(record_path)
+    await _force_remove(record_metadata_path)
+
+    if should_set:
+        if write_metadata:
+            async with aiofiles.open(record_metadata_path, mode='wb') as f:
+                await f.write(_json_dumps({
+                    'key': record['key'],
+                    'contentType': record.get('content_type') or 'unknown/no content type',
+                    'extension': record['extension'],
+                }).encode('utf-8'))
+
+        # Convert to bytes if string
+        if isinstance(record['value'], str):
+            record['value'] = record['value'].encode('utf-8')
+
+        async with aiofiles.open(record_path, mode='wb') as f:
+            await f.write(record['value'])
+
+
+async def _update_request_queue_item(
+    *,
+    request_id: str,
+    request: Dict,
+    entity_directory: str,
+    persist_storage: bool,
+) -> None:
+    # Skip writing files to the disk if the client has the option set to false
+    if not persist_storage:
+        return
+
+    # Ensure the directory for the entity exists
+    await makedirs(entity_directory, exist_ok=True)
+
+    # Write the request to the file
+    file_path = os.path.join(entity_directory, f'{request_id}.json')
+    async with aiofiles.open(file_path, mode='wb') as f:
+        await f.write(_json_dumps(request).encode('utf-8'))
+
+
+async def _delete_request(*, request_id: str, entity_directory: str) -> None:
+    # Ensure the directory for the entity exists
+    await makedirs(entity_directory, exist_ok=True)
+
+    file_path = os.path.join(entity_directory, f'{request_id}.json')
+    await _force_remove(file_path)