Skip to content

Commit

Permalink
Use Pydantic instead of asdict()
Browse files Browse the repository at this point in the history
We will use Pydantic later on for JSON decoding, so use it for
encoding now.
  • Loading branch information
KyleFromNVIDIA committed Jun 24, 2024
1 parent 1b6156a commit 70b38a7
Show file tree
Hide file tree
Showing 4 changed files with 11 additions and 18 deletions.
3 changes: 3 additions & 0 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -43,13 +43,16 @@ repos:
pass_filenames: false
additional_dependencies:
- packaging
- pydantic
- repo: local
hooks:
- id: generate-json
name: generate-json
entry: ./ci/generate_json.py
language: python
pass_filenames: false
additional_dependencies:
- pydantic

default_language_version:
python: python3
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ classifiers = [
requires-python = ">=3.9"
dependencies = [
"packaging",
"pydantic",
]

[project.scripts]
Expand Down
22 changes: 5 additions & 17 deletions src/rapids_metadata/json.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,19 +13,15 @@
# limitations under the License.

import argparse
import dataclasses
import json
import os
import sys
from typing import Any, Union
from typing import Union

from pydantic import TypeAdapter

from . import all_metadata
from .metadata import (
RAPIDSMetadata,
RAPIDSPackage,
RAPIDSRepository,
RAPIDSVersion,
)
from .metadata import RAPIDSMetadata
from .rapids_version import get_rapids_version


Expand All @@ -34,13 +30,6 @@
]


class _RAPIDSMetadataEncoder(json.JSONEncoder):
def default(
self, o: Union[RAPIDSMetadata, RAPIDSPackage, RAPIDSRepository, RAPIDSVersion]
) -> dict[str, Any]:
return dataclasses.asdict(o)


def main(argv: Union[list[str], None] = None):
if argv is None:
argv = sys.argv[1:]
Expand Down Expand Up @@ -77,9 +66,8 @@ def main(argv: Union[list[str], None] = None):

def write_file(f):
json.dump(
metadata,
TypeAdapter(RAPIDSMetadata).dump_python(metadata),
f,
cls=_RAPIDSMetadataEncoder,
sort_keys=True,
separators=(",", ": ") if parsed.pretty else (",", ":"),
indent=" " if parsed.pretty else None,
Expand Down
3 changes: 2 additions & 1 deletion tests/test_json.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from unittest.mock import patch

import pytest
from pydantic import TypeAdapter
from rapids_metadata import json as rapids_json
from rapids_metadata.metadata import (
RAPIDSMetadata,
Expand Down Expand Up @@ -133,7 +134,7 @@ def set_cwd(cwd: os.PathLike) -> Generator:
],
)
def test_metadata_encoder(unencoded, encoded):
assert rapids_json._RAPIDSMetadataEncoder().default(unencoded) == encoded
assert TypeAdapter(type(unencoded)).dump_python(unencoded) == encoded


@pytest.mark.parametrize(
Expand Down

0 comments on commit 70b38a7

Please sign in to comment.