add utility functions for generating random data given a schema and also to anonymize a schema
MODIFY
docs/index.rst
docs/index.rst
@@ -89,6 +89,7 @@ Documentation
json_writer
schema
validation
+ utils
command_line_script
* :ref:`genindex`

ADD
docs/utils.rst
+fastavro.utils
+==============
+
+.. autofunction:: fastavro.utils.generate_one
+.. autofunction:: fastavro.utils.generate_many
+.. autofunction:: fastavro.utils.anonymize_schema

MODIFY
fastavro/_schema.pyi
fastavro/_schema.pyi
@@ -6,7 +6,11 @@ FINGERPRINT_ALGORITHMS: str
def load_schema(schema_path: str) -> Dict: ...
def parse_schema(
- schema: SchemaType, expand: bool, _write_hint: bool, _force: bool
+ schema: SchemaType,
+ expand: bool = ...,
+ _write_hint: bool = ...,
+ _force: bool = ...,
+ _named_schemas: Dict = ...,
) -> SchemaType: ...
def fullname(schema: Dict) -> str: ...
def expand_schema(schema: Dict) -> Dict: ...

MODIFY
fastavro/_write.pyi
fastavro/_write.pyi
@@ -39,7 +39,7 @@ class Writer(GenericWriter):
): ...
def dump(self) -> None: ...
def write(self, record: AvroMessage) -> None: ...
- def write_block(self, block) -> None: ...
+ def write_block(self, block) -> None: ... # type: ignore # Should be a read.Block
def flush(self) -> None: ...
def schemaless_writer(fo: IO, schema: Dict, record: Dict) -> None: ...

RENAME
fastavro/types.pyi
fastavro/types.py
@@ -12,3 +12,5 @@ AvroMessage = Union[
List, # 'array'
Dict, # 'map' and 'record'
]
+Schema = Union[str, List, Dict]
+NamedSchemas = Dict[str, Dict]

ADD
fastavro/utils.py
+from hashlib import md5
+import random
+from string import ascii_letters
+from typing import Any, Iterator, Dict, List, cast
+
+from .const import INT_MIN_VALUE, INT_MAX_VALUE, LONG_MIN_VALUE, LONG_MAX_VALUE
+from .schema import extract_record_type, parse_schema
+from .types import Schema, NamedSchemas
+from ._schema_common import PRIMITIVES
+
+
+def _randbytes(num: int) -> bytes:
+ # TODO: Use random.randbytes when this library is Python 3.9+ only
+ return random.getrandbits(num * 8).to_bytes(num, "little")
+
+
+def _md5(string: str) -> str:
+ return md5(string.encode()).hexdigest()
+
+
+def _gen_utf8() -> str:
+ return "".join(random.choices(ascii_letters, k=10))
+
+
+def gen_data(schema: Schema, named_schemas: NamedSchemas, index: int) -> Any:
+
+ record_type = extract_record_type(schema)
+
+ if record_type == "null":
+ return None
+ elif record_type == "string":
+ return _gen_utf8()
+ elif record_type == "int":
+ return random.randint(INT_MIN_VALUE, INT_MAX_VALUE)
+ elif record_type == "long":
+ return random.randint(LONG_MIN_VALUE, LONG_MAX_VALUE)
+ elif record_type == "float":
+ return random.random()
+ elif record_type == "double":
+ return random.random()
+ elif record_type == "boolean":
+ return index % 2 == 0
+ elif record_type == "bytes":
+ return _randbytes(10)
+ elif record_type == "fixed":
+ fixed_schema = cast(Dict[str, Any], schema)
+ return _randbytes(fixed_schema["size"])
+ elif record_type == "enum":
+ enum_schema = cast(Dict[str, Any], schema)
+ real_index = index % len(enum_schema["symbols"])
+ return enum_schema["symbols"][real_index]
+ elif record_type == "array":
+ array_schema = cast(Dict[str, Schema], schema)
+ return [
+ gen_data(array_schema["items"], named_schemas, index) for _ in range(10)
+ ]
+ elif record_type == "map":
+ map_schema = cast(Dict[str, Schema], schema)
+ return {
+ _gen_utf8(): gen_data(map_schema["values"], named_schemas, index)
+ for _ in range(10)
+ }
+ elif record_type == "union" or record_type == "error_union":
+ union_schema = cast(List[Schema], schema)
+ real_index = index % len(union_schema)
+ return gen_data(union_schema[real_index], named_schemas, index)
+ elif record_type == "record" or record_type == "error":
+ record_schema = cast(Dict[str, Any], schema)
+ return {
+ field["name"]: gen_data(field["type"], named_schemas, index)
+ for field in record_schema["fields"]
+ }
+ else:
+ named_schema = cast(str, schema)
+ return gen_data(named_schemas[named_schema], named_schemas, index)
+
+
+def generate_one(schema: Schema) -> Any:
+ """
+ Returns a single instance of arbitrary data that conforms to the schema.
+
+ Parameters
+ ----------
+ schema: dict, list, string
+ Schema
+
+
+ Example::
+
+ from fastavro import schemaless_writer
+ from fastavro.utils import generate_one
+
+ schema = {
+ 'doc': 'A weather reading.',
+ 'name': 'Weather',
+ 'namespace': 'test',
+ 'type': 'record',
+ 'fields': [
+ {'name': 'station', 'type': 'string'},
+ {'name': 'time', 'type': 'long'},
+ {'name': 'temp', 'type': 'int'},
+ ],
+ }
+
+ with open('weather.avro', 'wb') as out:
+ schemaless_writer(out, schema, generate_one(schema))
+ """
+ return next(generate_many(schema, 1))
+
+
+def generate_many(schema: Schema, count: int) -> Iterator[Any]:
+ """
+ A generator that yields arbitrary data that conforms to the schema. It will
+ yield a number of data structures equal to what is given in the count
+
+ Parameters
+ ----------
+ schema: dict, list, string
+ Schema
+ count: int
+ Number of objects to generate
+
+
+ Example::
+
+ from fastavro import writer
+ from fastavro.utils import generate_data
+
+ schema = {
+ 'doc': 'A weather reading.',
+ 'name': 'Weather',
+ 'namespace': 'test',
+ 'type': 'record',
+ 'fields': [
+ {'name': 'station', 'type': 'string'},
+ {'name': 'time', 'type': 'long'},
+ {'name': 'temp', 'type': 'int'},
+ ],
+ }
+
+ with open('weather.avro', 'wb') as out:
+ writer(out, schema, generate_data(schema, 5))
+ """
+ named_schemas: NamedSchemas = {}
+ parsed_schema = parse_schema(schema, _named_schemas=named_schemas)
+ for index in range(count):
+ yield gen_data(parsed_schema, named_schemas, index)
+
+
+def anonymize_schema(schema: Schema) -> Schema:
+ """Returns an anonymized schema
+
+ Parameters
+ ----------
+ schema: dict, list, string
+ Schema
+
+
+ Example::
+
+ from fastavro.utils import anonymize_schema
+
+ anonymized_schema = anonymize_schema(original_schema)
+ """
+ named_schemas: NamedSchemas = {}
+ parsed_schema = parse_schema(schema, _named_schemas=named_schemas)
+ return _anonymize_schema(parsed_schema, named_schemas)
+
+
+def _anonymize_schema(schema: Schema, named_schemas: NamedSchemas) -> Schema:
+ # union schemas
+ if isinstance(schema, list):
+ return [_anonymize_schema(s, named_schemas) for s in schema]
+
+ # string schemas; this could be either a named schema or a primitive type
+ elif not isinstance(schema, dict):
+ if schema in PRIMITIVES:
+ return schema
+ else:
+ return _md5(schema)
+
+ else:
+ # Remaining valid schemas must be dict types
+ schema_type = schema["type"]
+
+ parsed_schema = {}
+ parsed_schema["type"] = schema_type
+
+ if "doc" in schema:
+ parsed_schema["doc"] = _md5(schema["doc"])
+
+ if schema_type == "array":
+ parsed_schema["items"] = _anonymize_schema(schema["items"], named_schemas)
+
+ elif schema_type == "map":
+ parsed_schema["values"] = _anonymize_schema(schema["values"], named_schemas)
+
+ elif schema_type == "enum":
+ parsed_schema["name"] = _md5(schema["name"])
+ parsed_schema["symbols"] = [_md5(symbol) for symbol in schema["symbols"]]
+
+ elif schema_type == "fixed":
+ parsed_schema["name"] = _md5(schema["name"])
+ parsed_schema["size"] = schema["size"]
+
+ elif schema_type == "record" or schema_type == "error":
+ # records
+ parsed_schema["name"] = _md5(schema["name"])
+ parsed_schema["fields"] = [
+ anonymize_field(field, named_schemas) for field in schema["fields"]
+ ]
+
+ elif schema_type in PRIMITIVES:
+ parsed_schema["type"] = schema_type
+
+ return parsed_schema
+
+
+def anonymize_field(
+ field: Dict[str, Any], named_schemas: NamedSchemas
+) -> Dict[str, Any]:
+ parsed_field: Dict[str, Any] = {}
+
+ if "doc" in field:
+ parsed_field["doc"] = _md5(field["doc"])
+ if "aliases" in field:
+ parsed_field["aliases"] = [_md5(alias) for alias in field["aliases"]]
+ if "default" in field:
+ parsed_field["default"] = field["default"]
+
+ # TODO: Defaults for enums should be hashed. Maybe others too?
+
+ parsed_field["name"] = _md5(field["name"])
+ parsed_field["type"] = _anonymize_schema(field["type"], named_schemas)
+
+ return parsed_field

MODIFY
run-tests.sh
run-tests.sh
@@ -34,6 +34,7 @@ EOF
if [[ "$RUN_MYPY" = "yes" ]]; then
echo "running mypy"
mypy ./fastavro
+ mypy --strict --follow-imports=silent --allow-any-generics fastavro/utils.py
else
echo "skipping mypy"
fi

MODIFY
tests/conftest.py
tests/conftest.py
import pytest
+import random
import time
from datetime import timezone, datetime, timedelta
+SEED = time.time()
+random.seed(SEED)
+
+
+def pytest_report_header(config):
+ return f"SEED is {SEED}"
+
@pytest.fixture(scope="function", autouse=True)
def look_for_cython_error(capfd):

MODIFY
tests/test_str_py3.py
tests/test_str_py3.py
from io import BytesIO
from os import SEEK_SET
-from random import choice, seed
+from random import choice
from string import ascii_uppercase, digits
import fastavro
@@ -10,8 +10,6 @@ def test_str_py3():
letters = ascii_uppercase + digits
id_size = 100
- seed("str_py3") # Repeatable results
-
def gen_id():
return "".join(choice(letters) for _ in range(id_size))

ADD
tests/test_utils.py
+import random
+from io import BytesIO
+from fastavro import schemaless_writer
+from fastavro.utils import generate_one, generate_many, anonymize_schema
+
+
+def test_generate():
+ schema = {
+ "type": "record",
+ "name": "Test",
+ "namespace": "test",
+ "fields": [
+ {"name": "null", "type": "null"},
+ {"name": "boolean", "type": "boolean"},
+ {"name": "string", "type": "string"},
+ {"name": "bytes", "type": "bytes"},
+ {"name": "int", "type": "int"},
+ {"name": "long", "type": "long"},
+ {"name": "float", "type": "float"},
+ {"name": "double", "type": "double"},
+ {
+ "name": "fixed",
+ "type": {"type": "fixed", "name": "fixed_field", "size": 5},
+ },
+ {
+ "name": "union",
+ "type": [
+ "null",
+ "int",
+ {
+ "type": "record",
+ "name": "union_record",
+ "fields": [{"name": "union_record_field", "type": "string"}],
+ },
+ ],
+ },
+ {
+ "name": "enum",
+ "type": {
+ "type": "enum",
+ "name": "enum_field",
+ "symbols": ["FOO", "BAR"],
+ },
+ },
+ {"name": "array", "type": {"type": "array", "items": "string"}},
+ {"name": "map", "type": {"type": "map", "values": "int"}},
+ {
+ "name": "record",
+ "type": {
+ "type": "record",
+ "name": "subrecord",
+ "fields": [{"name": "sub_int", "type": "int"}],
+ },
+ },
+ {"name": "named_type", "type": "subrecord"},
+ ],
+ }
+
+ count = 10
+
+ # Use list() to exhaust the generator)
+ assert len(list(generate_many(schema, count))) == count
+
+
+def test_anonymize():
+ schema = {
+ "type": "record",
+ "name": "Test",
+ "namespace": "test",
+ "doc": "this is a record",
+ "fields": [
+ {"name": "null", "type": "null"},
+ {"name": "boolean", "type": "boolean"},
+ {"name": "string", "type": "string", "default": "foo"},
+ {"name": "bytes", "type": "bytes", "aliases": ["alias_field"]},
+ {"name": "int", "type": "int", "doc": "doc"},
+ {"name": "long", "type": "long"},
+ {"name": "float", "type": "float"},
+ {"name": "double", "type": "double"},
+ {
+ "name": "fixed",
+ "type": {"type": "fixed", "name": "fixed_field", "size": 5},
+ },
+ {
+ "name": "union",
+ "type": [
+ "null",
+ "int",
+ {
+ "type": "record",
+ "name": "union_record",
+ "fields": [{"name": "union_record_field", "type": "string"}],
+ },
+ ],
+ },
+ {"name": "array", "type": {"type": "array", "items": "string"}},
+ {"name": "map", "type": {"type": "map", "values": "int"}},
+ {
+ "name": "record",
+ "type": {
+ "type": "record",
+ "name": "subrecord",
+ "fields": [{"name": "sub_int", "type": "int"}],
+ },
+ },
+ {"name": "named_type", "type": "subrecord"},
+ {"name": "other_int", "type": {"type": "int"}},
+ ],
+ }
+
+ anonymous_schema = anonymize_schema(schema)
+
+ # Maintain random state so that other tests continue to be based off the
+ # main starting seed
+ seed_state = random.getstate()
+ random.seed(1)
+ record = generate_one(schema)
+
+ random.seed(1)
+ anonymous_record = generate_one(anonymous_schema)
+ random.setstate(seed_state)
+
+ bio1 = BytesIO()
+ schemaless_writer(bio1, schema, record)
+
+ bio2 = BytesIO()
+ schemaless_writer(bio2, anonymous_schema, anonymous_record)
+
+ assert bio1.getvalue() == bio2.getvalue()
+
+
+def test_enum_symbols_get_anonymized():
+ schema = {
+ "type": "enum",
+ "name": "enum_field",
+ "symbols": ["FOO", "BAR"],
+ }
+
+ anonymous_schema = anonymize_schema(schema)
+
+ assert anonymous_schema["symbols"] != schema["symbols"]

add utility functions for generating random data given a schema and also to anonymize a schema