fix: move PredictionType to robotoff/types.py

legacy/ocr_category/category.py

@@ -4,7 +4,8 @@
 
 from robotoff.prediction.category.prediction_from_ocr.constants import LIST_CATEGORIES
 from robotoff.prediction.category.prediction_from_ocr.predictor import Predictor
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 
 from .dataclass import OCRResult, get_text
 

robotoff/app/api.py

@@ -54,10 +54,10 @@
 from robotoff.prediction.category import predict_category
 from robotoff.prediction.object_detection import ObjectDetectionModelRegistry
 from robotoff.prediction.ocr.dataclass import OCRParsingException
-from robotoff.prediction.types import PredictionType
 from robotoff.products import get_product_dataset_etag
 from robotoff.spellcheck import SPELLCHECKERS, Spellchecker
 from robotoff.taxonomy import is_prefixed_value, match_taxonomized_value
+from robotoff.types import PredictionType
 from robotoff.utils import get_image_from_url, get_logger, http_session
 from robotoff.utils.es import get_es_client
 from robotoff.utils.i18n import TranslationStore

robotoff/cli/insights.py

@@ -13,7 +13,8 @@
 from robotoff.off import get_barcode_from_path
 from robotoff.prediction.ocr import OCRResult, extract_predictions
 from robotoff.prediction.ocr.core import ocr_content_iter
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 from robotoff.utils import get_logger, gzip_jsonl_iter, jsonl_iter
 
 logger = get_logger(__name__)

robotoff/cli/main.py

@@ -5,7 +5,7 @@
 
 import typer
 
-from robotoff.types import ObjectDetectionModel, WorkerQueue
+from robotoff.types import ObjectDetectionModel, PredictionType, WorkerQueue
 
 app = typer.Typer()
 
@@ -92,7 +92,6 @@ def generate_ocr_predictions(
 ) -> None:
     """Generate OCR predictions of the requested type."""
     from robotoff.cli import insights
-    from robotoff.prediction.types import PredictionType
     from robotoff.utils import get_logger
 
     get_logger()
@@ -187,7 +186,6 @@ def import_insights(
     from robotoff.cli.insights import generate_from_ocr_archive, insights_iter
     from robotoff.insights import importer
     from robotoff.models import db
-    from robotoff.prediction.types import PredictionType
     from robotoff.utils import get_logger
 
     logger = get_logger()

robotoff/insights/extraction.py

@@ -11,8 +11,8 @@
     ObjectDetectionModelRegistry,
 )
 from robotoff.prediction.ocr.core import get_ocr_result
-from robotoff.prediction.types import Prediction, PredictionType
-from robotoff.types import ObjectDetectionModel
+from robotoff.prediction.types import Prediction
+from robotoff.types import ObjectDetectionModel, PredictionType
 from robotoff.utils import get_logger, http_session
 
 logger = get_logger(__name__)

robotoff/insights/importer.py

@@ -15,7 +15,7 @@
 from robotoff.models import Prediction as PredictionModel
 from robotoff.models import ProductInsight, batch_insert
 from robotoff.off import get_server_type
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
 from robotoff.products import (
     DBProductStore,
     Product,
@@ -30,6 +30,7 @@
     get_taxonomy,
     match_taxonomized_value,
 )
+from robotoff.types import PredictionType
 from robotoff.utils import get_logger, text_file_iter
 from robotoff.utils.cache import CachedStore
 

robotoff/logos.py

@@ -8,8 +8,9 @@
 from robotoff.insights.importer import import_insights
 from robotoff.logo_label_type import LogoLabelType
 from robotoff.models import ImageModel, LogoAnnotation, LogoConfidenceThreshold
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
 from robotoff.slack import NotifierFactory
+from robotoff.types import PredictionType
 from robotoff.utils import get_logger, http_session
 from robotoff.utils.types import JSONType
 

robotoff/prediction/category/matcher.py

@@ -8,9 +8,10 @@
 from flashtext import KeywordProcessor
 
 from robotoff import settings
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
 from robotoff.products import ProductDataset
 from robotoff.taxonomy import TaxonomyType, get_taxonomy
+from robotoff.types import PredictionType
 from robotoff.utils import dump_json, get_logger, load_json
 from robotoff.utils.text import (
     get_lemmatizing_nlp,

robotoff/prediction/category/neural/category_classifier.py

@@ -1,8 +1,9 @@
 from typing import Optional
 
 from robotoff import settings
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
 from robotoff.taxonomy import Taxonomy
+from robotoff.types import PredictionType
 from robotoff.utils import http_session
 
 

robotoff/prediction/ocr/brand.py

@@ -5,7 +5,8 @@
 
 from robotoff import settings
 from robotoff.brands import BRAND_BLACKLIST_STORE, keep_brand_from_taxonomy
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 from robotoff.utils import get_logger, text_file_iter
 from robotoff.utils.text import get_tag
 

robotoff/prediction/ocr/category.py

@@ -2,8 +2,9 @@
 from typing import Optional, Union
 
 from robotoff.off import normalize_tag
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
 from robotoff.taxonomy import get_taxonomy
+from robotoff.types import PredictionType
 from robotoff.utils import get_logger
 
 from .dataclass import OCRField, OCRRegex, OCRResult, get_text

robotoff/prediction/ocr/core.py

@@ -5,8 +5,9 @@
 import requests
 
 from robotoff.off import generate_json_ocr_url, get_barcode_from_path, split_barcode
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
 from robotoff.settings import BaseURLProvider
+from robotoff.types import PredictionType
 from robotoff.utils import get_logger, http_session, jsonl_iter, jsonl_iter_fp
 from robotoff.utils.types import JSONType
 

robotoff/prediction/ocr/expiration_date.py

@@ -3,7 +3,8 @@
 import re
 from typing import Optional, Union
 
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 
 from .dataclass import OCRField, OCRRegex, OCRResult, get_text
 

robotoff/prediction/ocr/image_flag.py

@@ -3,7 +3,8 @@
 from flashtext import KeywordProcessor
 
 from robotoff import settings
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 from robotoff.utils import text_file_iter
 
 from .dataclass import OCRResult, SafeSearchAnnotationLikelihood, get_text

robotoff/prediction/ocr/image_lang.py

@@ -1,6 +1,7 @@
 from typing import Optional, Union
 
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 
 from .dataclass import OCRResult
 

robotoff/prediction/ocr/image_orientation.py

@@ -1,6 +1,7 @@
 from typing import Union
 
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 
 from .dataclass import ImageOrientation, OCRResult
 

robotoff/prediction/ocr/label.py

@@ -4,7 +4,8 @@
 from flashtext import KeywordProcessor
 
 from robotoff import settings
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 from robotoff.utils import get_logger, text_file_iter
 from robotoff.utils.cache import CachedStore
 

robotoff/prediction/ocr/location.py

@@ -8,7 +8,8 @@
 from flashtext import KeywordProcessor
 
 from robotoff import settings
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 from robotoff.utils import get_logger
 from robotoff.utils.cache import CachedStore
 from robotoff.utils.text import strip_accents_ascii

robotoff/prediction/ocr/nutrient.py

@@ -1,7 +1,8 @@
 import re
 from typing import Union
 
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 from robotoff.utils.types import JSONType
 
 from .dataclass import OCRField, OCRRegex, OCRResult, get_text

robotoff/prediction/ocr/packager_code.py

@@ -4,7 +4,8 @@
 from flashtext import KeywordProcessor
 
 from robotoff import settings
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 from robotoff.utils import text_file_iter
 from robotoff.utils.cache import CachedStore
 

robotoff/prediction/ocr/product_weight.py

@@ -4,7 +4,8 @@
 
 import pint
 
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 from robotoff.utils import get_logger
 
 from .dataclass import OCRField, OCRRegex, OCRResult, get_text

robotoff/prediction/ocr/store.py

@@ -2,7 +2,8 @@
 from typing import Union
 
 from robotoff import settings
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 from robotoff.utils import text_file_iter
 
 from .dataclass import OCRField, OCRRegex, OCRResult, get_text

robotoff/prediction/ocr/trace.py

@@ -2,7 +2,8 @@
 from typing import Optional, Union
 
 from robotoff import settings
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
+from robotoff.types import PredictionType
 from robotoff.utils import text_file_iter
 from robotoff.utils.cache import CachedStore
 

robotoff/prediction/types.py

@@ -1,35 +1,8 @@
 import dataclasses
 import datetime
-from enum import Enum, unique
 from typing import Any, Optional
 
-
-@unique
-class PredictionType(str, Enum):
-    """PredictionType defines the type of the prediction.
-
-    See `InsightType` documentation for further information about each type.
-    """
-
-    ingredient_spellcheck = "ingredient_spellcheck"
-    packager_code = "packager_code"
-    label = "label"
-    category = "category"
-    image_flag = "image_flag"
-    product_weight = "product_weight"
-    expiration_date = "expiration_date"
-    brand = "brand"
-    image_orientation = "image_orientation"
-    store = "store"
-    nutrient = "nutrient"
-    trace = "trace"
-    packaging_text = "packaging_text"
-    packaging = "packaging"
-    location = "location"
-    nutrient_mention = "nutrient_mention"
-    image_lang = "image_lang"
-    nutrition_image = "nutrition_image"
-    nutrition_table_structure = "nutrition_table_structure"
+from robotoff.types import PredictionType
 
 
 @dataclasses.dataclass

robotoff/scheduler/latent.py

@@ -6,7 +6,6 @@
 from robotoff.insights.dataclass import InsightType
 from robotoff.models import Prediction, ProductInsight, with_db
 from robotoff.off import get_server_type
-from robotoff.prediction.types import PredictionType
 from robotoff.products import (
     DBProductStore,
     get_image_id,
@@ -15,6 +14,7 @@
     is_nutrition_image,
     is_valid_image,
 )
+from robotoff.types import PredictionType
 from robotoff.utils import get_logger
 from robotoff.utils.types import JSONType
 

robotoff/spellcheck/__init__.py

@@ -1,14 +1,15 @@
 from typing import Iterable, Optional, Union
 
 from robotoff import settings
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
 from robotoff.products import ProductDataset
 from robotoff.spellcheck.base_spellchecker import BaseSpellchecker
 from robotoff.spellcheck.elasticsearch import ElasticSearchSpellchecker
 from robotoff.spellcheck.items import SpellcheckItem
 from robotoff.spellcheck.patterns import PatternsSpellchecker
 from robotoff.spellcheck.percentages import PercentagesSpellchecker
 from robotoff.spellcheck.vocabulary import VocabularySpellchecker
+from robotoff.types import PredictionType
 from robotoff.utils.types import JSONType
 
 SPELLCHECKERS = {

robotoff/types.py

@@ -10,3 +10,31 @@ class ObjectDetectionModel(enum.Enum):
     nutriscore = "nutriscore"
     universal_logo_detector = "universal-logo-detector"
     nutrition_table = "nutrition-table"
+
+
+@enum.unique
+class PredictionType(str, enum.Enum):
+    """PredictionType defines the type of the prediction.
+
+    See `InsightType` documentation for further information about each type.
+    """
+
+    ingredient_spellcheck = "ingredient_spellcheck"
+    packager_code = "packager_code"
+    label = "label"
+    category = "category"
+    image_flag = "image_flag"
+    product_weight = "product_weight"
+    expiration_date = "expiration_date"
+    brand = "brand"
+    image_orientation = "image_orientation"
+    store = "store"
+    nutrient = "nutrient"
+    trace = "trace"
+    packaging_text = "packaging_text"
+    packaging = "packaging"
+    location = "location"
+    nutrient_mention = "nutrient_mention"
+    image_lang = "image_lang"
+    nutrition_image = "nutrition_image"
+    nutrition_table_structure = "nutrition_table_structure"

robotoff/workers/tasks/import_image.py

@@ -27,9 +27,9 @@
 )
 from robotoff.off import get_server_type, get_source_from_url
 from robotoff.prediction.object_detection import ObjectDetectionModel
-from robotoff.prediction.types import PredictionType
 from robotoff.products import Product, get_product_store
 from robotoff.slack import NotifierFactory
+from robotoff.types import PredictionType
 from robotoff.utils import get_image_from_url, get_logger, http_session
 from robotoff.workers.queues import enqueue_job, high_queue
 

tests/integration/insights/test_category_import.py

@@ -3,8 +3,9 @@
 from robotoff import settings
 from robotoff.insights.importer import import_insights
 from robotoff.models import ProductInsight
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
 from robotoff.products import Product
+from robotoff.types import PredictionType
 
 from ..models_utils import PredictionFactory, ProductInsightFactory, clean_db
 

tests/unit/insights/test_importer.py

@@ -23,9 +23,10 @@
     sort_predictions,
 )
 from robotoff.models import ProductInsight
-from robotoff.prediction.types import Prediction, PredictionType
+from robotoff.prediction.types import Prediction
 from robotoff.products import Product
 from robotoff.taxonomy import get_taxonomy
+from robotoff.types import PredictionType
 
 DEFAULT_BARCODE = "3760094310634"
 DEFAULT_SERVER_DOMAIN = "api.openfoodfacts.org"