Skip to content
Snippets Groups Projects
number_replacer.py 1.28 KiB
Newer Older
Michał Pogoda's avatar
Michał Pogoda committed
from typing import List, Tuple
from src.detections import (
    Detection,
    NumberDetection,
)
from src.string_replacements import replace_and_update
from src.replacers.interface import ReplacerInterface
import random
import string

Michał Pogoda's avatar
Michał Pogoda committed
def randomize_digits_in_text(text: str) -> str:
    result = ""
Michał Pogoda's avatar
Michał Pogoda committed
    for c in text:
        if c.isdigit():
            result += random.choice(string.digits)
        else:
            result += c
Michał Pogoda's avatar
Michał Pogoda committed
    return result

Michał Pogoda's avatar
Michał Pogoda committed
class NumberReplacer(ReplacerInterface):
    def __init__(self):
        pass
Michał Pogoda's avatar
Michał Pogoda committed

    def replace(
        self, text: str, detections: List[Tuple[int, int, Detection]]
    ) -> Tuple[str, List[Tuple[int, int, Detection]]]:
Michał Pogoda's avatar
Michał Pogoda committed
        replacements = []
        not_processed = []
Michał Pogoda's avatar
Michał Pogoda committed
        already_replaced = dict()

        for item in detections:
            start, end, detection = item

            if isinstance(detection, NumberDetection):
                if text[start:end] not in already_replaced:
Michał Pogoda's avatar
Michał Pogoda committed
                    already_replaced[text[start:end]] = randomize_digits_in_text(
                        text[start:end]
                    )

Michał Pogoda's avatar
Michał Pogoda committed
                replacements.append((start, end, already_replaced[text[start:end]]))
            else:
                not_processed.append(item)
Michał Pogoda's avatar
Michał Pogoda committed

        return replace_and_update(text, replacements, not_processed)