refactor: add type hints for user_id and visualization server args (#1987)
<!-- .github/pull_request_template.md --> ## Description Resolves #1986 I addressed all `ANN001` errors in `cognee/shared/utils.py`. Updated functions: 1. send_telemetry 2. start_visualization_server 3. _sanitize_nested_properties 4. embed_logo While fixing these errors, i've noticed that the `send_telemetry` function lacked a type hint for `user_id`. After analyzing the `User` models and usage patterns in the codebase, I found that `user_id` is not strictly a `str` but can also be a `uuid.UUID` object. Therefore, I updated the type hint to `Union[str, uuid.UUID]` (importing `uuid` and `typing.Union`) to accurately reflect the data structure and improve type safety. ## Acceptance Criteria * [x] The code passes static analysis (`ruff`) without `ANN001` errors in `cognee/shared/utils.py`. * [x] Correct imports (`uuid`, `Union`) are added and sorted. [Check Steps] 1. Run 'uv run ruff check cognee/shared/utils.py --select ANN001' 5. Expected result: No errors found. ## Type of Change <!-- Please check the relevant option --> - [ ] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [x] Code refactoring - [ ] Performance improvement - [ ] Other (please specify): ## Screenshots/Videos (if applicable) <!-- Add screenshots or videos to help explain your changes --> ## Pre-submission Checklist <!-- Please check all boxes that apply before submitting your PR --> - [x] **I have tested my changes thoroughly before submitting this PR** - [x] **This PR contains minimal changes necessary to address the issue/feature** - [x] My code follows the project's coding standards and style guidelines - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] I have added necessary documentation (if applicable) - [x] All new and existing tests pass - [x] I have searched existing PRs to ensure this change hasn't been submitted already - [x] I have linked any relevant issues in the description - [x] My commits have clear and descriptive messages ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Refactor** * Improved type annotations across telemetry and sanitization utilities for safer handling of IDs and nested properties. * Ensured additional properties are sanitized before telemetry is sent. * Added explicit type hints for visualization startup and logo embedding parameters for clearer IDE support. This release contains no user-facing changes. <sub>✏️ Tip: You can customize this high-level summary in your review settings.</sub> <!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
commit
861c5e33da
1 changed files with 12 additions and 5 deletions
|
|
@ -8,7 +8,8 @@ import http.server
|
|||
import socketserver
|
||||
from threading import Thread
|
||||
import pathlib
|
||||
from uuid import uuid4, uuid5, NAMESPACE_OID
|
||||
from typing import Union, Any, Dict, List
|
||||
from uuid import uuid4, uuid5, NAMESPACE_OID, UUID
|
||||
|
||||
from cognee.base_config import get_base_config
|
||||
from cognee.shared.logging_utils import get_logger
|
||||
|
|
@ -58,7 +59,7 @@ def get_anonymous_id():
|
|||
return anonymous_id
|
||||
|
||||
|
||||
def _sanitize_nested_properties(obj, property_names: list[str]):
|
||||
def _sanitize_nested_properties(obj: Any, property_names: list[str]) -> Any:
|
||||
"""
|
||||
Recursively replaces any property whose key matches one of `property_names`
|
||||
(e.g., ['url', 'path']) in a nested dict or list with a uuid5 hash
|
||||
|
|
@ -78,7 +79,9 @@ def _sanitize_nested_properties(obj, property_names: list[str]):
|
|||
return obj
|
||||
|
||||
|
||||
def send_telemetry(event_name: str, user_id, additional_properties: dict = {}):
|
||||
def send_telemetry(event_name: str, user_id: Union[str, UUID], additional_properties: dict = {}):
|
||||
if additional_properties is None:
|
||||
additional_properties = {}
|
||||
if os.getenv("TELEMETRY_DISABLED"):
|
||||
return
|
||||
|
||||
|
|
@ -108,7 +111,7 @@ def send_telemetry(event_name: str, user_id, additional_properties: dict = {}):
|
|||
print(f"Error sending telemetry through proxy: {response.status_code}")
|
||||
|
||||
|
||||
def embed_logo(p, layout_scale, logo_alpha, position):
|
||||
def embed_logo(p: Any, layout_scale: float, logo_alpha: float, position: str):
|
||||
"""
|
||||
Embed a logo into the graph visualization as a watermark.
|
||||
"""
|
||||
|
|
@ -138,7 +141,11 @@ def embed_logo(p, layout_scale, logo_alpha, position):
|
|||
|
||||
|
||||
def start_visualization_server(
|
||||
host="0.0.0.0", port=8001, handler_class=http.server.SimpleHTTPRequestHandler
|
||||
host: str = "0.0.0.0",
|
||||
port: int = 8001,
|
||||
handler_class: type[
|
||||
http.server.SimpleHTTPRequestHandler
|
||||
] = http.server.SimpleHTTPRequestHandler,
|
||||
):
|
||||
"""
|
||||
Spin up a simple HTTP server in a background thread to serve files.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue