refactor classes into separate files

Author: nebfield

pgscatalog.core/src/pgscatalog/core/__init__.py

@@ -10,7 +10,6 @@
     ScoringFile,
     NormalisedScoringFile,
     ScoreVariant,
-    EffectType,
     GenomeBuild,
     TargetVariants,
     TargetVariant,
@@ -37,6 +36,8 @@
     relabel_write,
     effect_type_keyfunc,
     chrom_keyfunc,
+    EffectType,
+    Allele,
 )
 
 log_fmt = "%(name)s: %(asctime)s %(levelname)-8s %(message)s"
@@ -73,12 +74,13 @@
     "TargetVariants",
     "TargetType",
     "NormalisedScoringFile",
-    "EffectType",
     "RelabelArgs",
     "relabel",
     "relabel_write",
     "effect_type_keyfunc",
     "chrom_keyfunc",
+    "EffectType",
+    "Allele",
 ]
 
 __version__ = importlib.metadata.version("pgscatalog.core")

pgscatalog.core/src/pgscatalog/core/lib/__init__.py

@@ -1,8 +1,11 @@
 from ._config import Config
+from .allele import Allele
+from .effecttype import EffectType
+from .genomebuild import GenomeBuild
 from .catalogapi import ScoreQueryResult, CatalogQuery, CatalogCategory
 from .scorefiles import ScoringFiles, ScoringFile, NormalisedScoringFile
-from .scorevariant import ScoreVariant, EffectType
-from .genomebuild import GenomeBuild
+from .catalogscorevariant import CatalogScoreVariant
+from .scorevariant import ScoreVariant
 from .targetvariants import TargetVariants, TargetVariant, TargetType
 from ._relabel import RelabelArgs, relabel, relabel_write
 from ._sortpaths import effect_type_keyfunc, chrom_keyfunc
@@ -56,10 +59,12 @@
     "TargetVariants",
     "TargetType",
     "NormalisedScoringFile",
-    "EffectType",
     "RelabelArgs",
     "relabel",
     "relabel_write",
     "effect_type_keyfunc",
     "chrom_keyfunc",
+    "CatalogScoreVariant",
+    "EffectType",
+    "Allele",
 ]

pgscatalog.core/src/pgscatalog/core/lib/_normalise.py

@@ -9,7 +9,7 @@
 import pyliftover
 
 from .genomebuild import GenomeBuild
-from .scorevariant import Allele
+from .allele import Allele
 from .pgsexceptions import LiftoverError
 
 logger = logging.getLogger(__name__)

pgscatalog.core/src/pgscatalog/core/lib/allele.py

@@ -0,0 +1,58 @@
+from functools import cached_property
+from typing import ClassVar
+
+from pydantic import BaseModel, computed_field, model_serializer
+
+
+class Allele(BaseModel):
+    """A class that represents an allele found in PGS Catalog scoring files
+    >>> simple_ea = Allele(**{"allele": "A"})
+    >>> simple_ea
+    Allele(allele='A', is_snp=True)
+    >>> str(simple_ea)
+    'A'
+    >>> Allele(**{"allele": "AG"})
+    Allele(allele='AG', is_snp=True)
+    >>> hla_example = Allele(**{"allele": "+"})
+    >>> hla_example
+    Allele(allele='+', is_snp=False)
+
+    >>> Allele(allele="A")
+    Allele(allele='A', is_snp=True)
+
+    >>> Allele(allele="A/T").has_multiple_alleles
+    True
+    """
+
+    allele: str
+    _valid_snp_bases: ClassVar[frozenset[str]] = frozenset({"A", "C", "T", "G"})
+
+    @computed_field
+    @cached_property
+    def is_snp(self) -> bool:
+        """SNPs are the most common type of effect allele in PGS Catalog scoring
+        files. More complex effect alleles, like HLAs or APOE genes, often require
+        extra work to represent in genomes. Users should be warned about complex
+        effect alleles.
+        """
+        return not frozenset(self.allele) - self._valid_snp_bases
+
+    @cached_property
+    def has_multiple_alleles(self) -> bool:
+        return "/" in self.allele
+
+    @model_serializer(mode="plain", return_type=str)
+    def serialize(self):
+        """When dumping the model, flatten it to just return the allele as a string"""
+        return self.allele
+
+    def __str__(self):
+        return self.allele
+
+    def __eq__(self, other):
+        if isinstance(other, Allele):
+            return self.allele == other.allele
+        return False
+
+    def __hash__(self):
+        return hash(self.allele)