Add TokenPolicy class

Author: astrojarred

pycardano/easy.py

@@ -240,9 +240,143 @@ def ad_ada(self):
         return Ada(self.ada)
 
 
+@dataclass(unsafe_hash=True)
+class TokenPolicy:
+    name: str
+    policy: Optional[Union[NativeScript, dict]] = field(repr=False, default=None)
+    policy_dir: Optional[Union[str, Path]] = field(
+        repr=False, default=Path("./priv/policies")
+    )
+
+    def __post_init__(self):
+
+        # streamline inputs
+        if isinstance(self.policy_dir, str):
+            self.policy_dir = Path(self.policy_dir)
+
+        if not self.policy_dir.exists():
+            self.policy_dir.mkdir(parents=True, exist_ok=True)
+
+        # look for the policy
+        if Path(self.policy_dir / f"{self.name}.script").exists():
+            with open(
+                Path(self.policy_dir / f"{self.name}.script"), "r"
+            ) as policy_file:
+                self.policy = NativeScript.from_dict(json.load(policy_file))
+
+        elif isinstance(self.policy, dict):
+            self.policy = NativeScript.from_dict(self.policy)
+
+    @property
+    def policy_id(self):
+
+        if self.policy:
+            return str(self.policy.hash())
+
+    @property
+    def expiration_slot(self):
+        """Get the expiration slot for a simple minting policy,
+        like one generated by generate_minting_policy
+        """
+
+        if self.policy:
+            scripts = getattr(self.policy, "native_scripts", None)
+
+            if scripts:
+                for script in scripts:
+                    if script._TYPE == 5:
+                        return script.after
+
+    def get_expiration_timestamp(self, context: ChainContext):
+        """Get the expiration timestamp for a simple minting policy,
+        like one generated by generate_minting_policy
+        """
+
+        if self.expiration_slot:
+
+            seconds_diff = self.expiration_slot - context.last_block_slot
+
+            return datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
+                seconds=seconds_diff
+            )
+
+    def is_expired(self, context: ChainContext):
+        """Get the expiration timestamp for a simple minting policy,
+        like one generated by generate_minting_policy
+        """
+
+        if self.expiration_slot:
+
+            seconds_diff = self.expiration_slot - context.last_block_slot
+
+            return seconds_diff < 0
+
+    def generate_minting_policy(
+        self,
+        signers: Union["Wallet", Address, List["Wallet"], List[Address]],
+        expiration: Optional[Union[datetime.datetime, int]] = None,
+        context: Optional[ChainContext] = None,
+    ):
+
+        script_filepath = Path(self.policy_dir / f"{self.name}.script")
+
+        if script_filepath.exists() or self.policy:
+            raise FileExistsError(f"Policy named {self.name} already exists")
+
+        if isinstance(expiration, datetime.datetime) and not context:
+            raise AttributeError(
+                "If input expiration is provided as a datetime, please also provide a context."
+            )
+
+        # get pub key hashes
+        if not isinstance(signers, list):
+            signers = [signers]
+
+        pub_keys = [ScriptPubkey(self._get_pub_key_hash(signer)) for signer in signers]
+
+        # calculate when to lock
+        if expiration:
+            if isinstance(expiration, int):  # assume this is directly the block no.
+                must_before_slot = InvalidHereAfter(expiration)
+            elif isinstance(expiration, datetime.datetime):
+                if expiration.tzinfo:
+                    time_until_expiration = expiration - datetime.datetime.now(
+                        datetime.datetime.utc
+                    )
+                else:
+                    time_until_expiration = expiration - datetime.datetime.now()
+
+                last_block_slot = context.last_block_slot
+
+                must_before_slot = InvalidHereAfter(
+                    last_block_slot + int(time_until_expiration.total_seconds())
+                )
+
+            policy = ScriptAll(pub_keys + [must_before_slot])
+
+        else:
+            policy = ScriptAll(pub_keys)
+
+        # save policy to file
+        with open(script_filepath, "w") as script_file:
+            json.dump(policy.to_dict(), script_file, indent=4)
+
+        self.policy = policy
+
+    @staticmethod
+    def _get_pub_key_hash(signer=Union["Wallet", Address]):
+
+        if hasattr(signer, "verification_key"):
+            return signer.verification_key.hash()
+        elif isinstance(signer, Address):
+            return str(signer.payment_part)
+        else:
+            raise TypeError("Input signer must be of type Wallet or Address.")
+
+
 @dataclass(unsafe_hash=True)
 class Token:
-    policy: Union[NativeScript, str]
+    policy: Union[NativeScript, TokenPolicy]
     amount: int
     name: Optional[str] = field(default="")
     hex_name: Optional[str] = field(default="")