refactor FileStatus as class, add --mathlib4

mathlibtools/file_status.py

@@ -1,53 +1,60 @@
-from dataclasses import dataclass, field
-from typing import Optional, Set
+import re
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, Optional, Union
 
-@dataclass(frozen=True)
+import requests
+import yaml
+
+
+@dataclass()
 class FileStatus:
-  """Capture the file status descriptions in the wiki yaml.
-
-  `string_match` are the substrings that must be found, lowercase match
-  `color` is how the node should be colored if it has the status
-  """
-
-  color: str
-  prefix: Optional[str] = None
-  string_match: Set[str] = field(default_factory=set)
-  # colors from X11
-
-  @classmethod
-  def yes(cls) -> "FileStatus":
-    return cls("green", "Yes")
-
-  @classmethod
-  def pr(cls) -> "FileStatus":
-    return cls("lightskyblue", "No", {"pr"})
-
-  @classmethod
-  def wip(cls) -> "FileStatus":
-    return cls("lightskyblue", "No", {"wip"})
-
-  @classmethod
-  def no(cls) -> "FileStatus":
-    return cls("orange", "No")
-
-  # @classmethod
-  # def missing(cls) -> "FileStatus":
-  #   return cls("orchid1")
-
-  @classmethod
-  def ready(cls) -> "FileStatus":
-    return cls("turquoise1")
-
-  def matches(self, comment: str) -> bool:
-    return (self.prefix is None or comment.startswith(self.prefix)) and \
-        all(substring.lower() in comment.lower() for substring in self.string_match)
-
-  @classmethod
-  def assign(cls, comment: str) -> Optional["FileStatus"]:
-    for status in [cls.yes(), cls.pr(), cls.wip()]:
-      if status.matches(comment):
-        return status
-    # "No" but other comments
-    if cls.no().matches(comment) and len(comment) > 2:
-      return cls.no()
-    return None
+
+    ported: bool = False
+    mathlib4_pr: Optional[int] = None
+    mathlib3_hash: Optional[str] = None
+    comments: Optional[str] = None
+
+    @classmethod
+    def parse_old(cls, message: str) -> "FileStatus":
+        ported = False
+        mathlib4_pr: Optional[int] = None
+        mathlib3_hash: Optional[str] = None
+        if message.startswith("Yes"):
+            ported = True
+            if len(message.split()) > 2:
+                mathlib3_hash = message.split()[2]
+        if "mathlib4#" in message:
+            mathlib4_pr = int(re.findall(r"[0-9]+", message.replace("mathlib4#", ""))[0])
+        return cls(
+            ported=ported,
+            mathlib4_pr=mathlib4_pr,
+            mathlib3_hash=mathlib3_hash,
+        )
+
+    def asdict(self) -> Dict[str, Any]:
+        return {k: v for k, v in asdict(self).items() if v is not None}
+
+
+
+@dataclass
+class PortStatus:
+
+    file_statuses: Dict[str, FileStatus]
+
+    @classmethod
+    def old_yaml(cls, url: Optional[str] = None) -> Dict[str, str]:
+        if url is None:
+            url = "https://raw.githubusercontent.com/wiki/leanprover-community/mathlib/mathlib4-port-status.md"
+        def yaml_md_load(wikicontent: bytes):
+            return yaml.safe_load(wikicontent.replace(b"```", b""))
+
+        return yaml_md_load(requests.get(url).content)
+
+    @classmethod
+    def deserialize_old(cls, yaml: Optional[Dict[str, str]] = None) -> "PortStatus":
+        if yaml is None:
+            yaml = cls.old_yaml()
+        return cls(file_statuses={k: FileStatus.parse_old(v) for k, v in yaml.items()})
+
+    def serialize(self) -> Dict[str, Dict[str, Union[int, str, None]]]:
+        return yaml.dump({k: v.asdict() for k, v in self.file_statuses.items()})

mathlibtools/import_graph.py

@@ -101,7 +101,7 @@ def transitive_reduction(self) -> 'ImportGraph':
     def delete_ported(self) -> 'ImportGraph':
         """Delete all nodes marked as ported during port_status"""
         H = self.subgraph({node for node, attrs in self.nodes(data=True)
-                          if attrs.get("status") != FileStatus.yes()})
+                          if attrs.get("status") and attrs.get("status").ported})
         H.base_path = self.base_path
         return H
 

mathlibtools/leanproject.py

@@ -1,7 +1,7 @@
 import sys
 from pathlib import Path
 from datetime import datetime
-from typing import Tuple, Optional
+from typing import List, Tuple, Optional
 from getpass import getpass
 
 from git.exc import GitCommandError # type: ignore
@@ -302,6 +302,8 @@ def global_upgrade() -> None:
               help='Color by mathlib4 porting status')
 @click.option('--port-status-url', default=None,
               help='URL of yaml with mathlib4 port status')
+@click.option('--mathlib4', default=None,
+              help='Local directory of mathlib4 repo')
 @click.option('--reduce', 'reduce', default=False, is_flag=True,
               help='Omit transitive imports.')
 @click.option('--show-unused', 'unused', default=False, is_flag=True,
@@ -313,6 +315,7 @@ def import_graph(
     exclude : bool,
     port_status: bool,
     port_status_url: Optional[str],
+    mathlib4: Optional[str],
     reduce: bool,
     unused: bool,
     output: str
@@ -334,7 +337,7 @@ def import_graph(
     if unused and to:
         project.show_unused(to)
     if port_status or port_status_url:
-        project.port_status(port_status_url)
+        project.port_status(port_status_url, mathlib4=None if mathlib4 is None else Path(mathlib4))
     if to and from_:
         G = graph.path(start=from_, end=to)
     elif to:
@@ -384,7 +387,7 @@ def port_progress(to: Optional[str]) -> None:
     print(f"| {'Longest unported chain:':<{W       }} | {longest_unported_path:>8}/{mathlib3_longest_path:<8} | ({progress_path:>3}% progress) |")
     print()
 
-    path = graph.longest_path()
+    path: List[str] = graph.longest_path()
     if path[-1] == "all":
         path = path[:-1]
     if not to:

mathlibtools/lib.py

@@ -27,7 +27,7 @@
                  GitCommandError, BadName, RemoteReference) # type: ignore
 from atomicwrites import atomic_write
 
-from mathlibtools.file_status import FileStatus
+from mathlibtools.file_status import FileStatus, PortStatus
 
 if TYPE_CHECKING:
     from mathlibtools.import_graph import ImportGraph
@@ -1019,34 +1019,37 @@ def pull(self, remote: str='origin') -> None:
                 self.run_echo(['leanpkg', 'configure'])
             self.get_mathlib_olean()
 
-    def port_status(self, url: Optional[str] = None) -> None:
+    def port_status(self, url: Optional[str] = None, mathlib4: Optional[Path] = None) -> None:
         """
         Color nodes on the graph based on port status. Done in place to the graph nodes.
 
         Args:
             url: md or yaml file with "file: label" content, by default, from the wiki
         """
-        if url is None:
-            url = 'https://raw.githubusercontent.com/wiki/leanprover-community/mathlib/mathlib4-port-status.md'
-        def yaml_md_load(wikicontent: bytes):
-            return yaml.safe_load(wikicontent.replace(b"```", b""))
 
-        port_labels: Dict[str, str] = yaml_md_load(requests.get(url).content)
+        def snake_to_camel(filename: str) -> str:
+            return re.sub('(?!^)([A-Z]+)', r'_\1', filename).lower()
 
-        for filename, status in port_labels.items():
-            if filename not in self.import_graph.nodes:
-                continue
-            node = self.import_graph.nodes[filename]
-            node_path = self.src_directory.joinpath(*filename.split(".")).with_suffix(".lean")
+        existing_files = set()
+        if mathlib4:
+            mathlib4 = mathlib4 / "Mathlib"
+            existing_files = {str.join(".",
+                map(snake_to_camel, leanfile.relative_to(mathlib4).with_suffix("").parts)).lower()
+                for leanfile in mathlib4.rglob("*") if leanfile.suffix == ".lean"}
+
+        port_status = PortStatus.deserialize_old()
+
+        for node_name, node in self.import_graph.nodes(data=True):
+            node["status"] = port_status.file_statuses.get(node_name, FileStatus())
+            node_path = self.src_directory.joinpath(*node_name.split(".")).with_suffix(".lean")
             with open(node_path, "r",encoding="utf-8",errors='ignore') as f:
                 node["nb_lines"] = sum(bl.count("\n") for bl in blocks(f))
-            node["status"] = FileStatus.assign(status)
         # somehow missing from yaml
         # for node_name, node in self.import_graph.nodes(data=True):
         #     if node_name not in port_labels:
         #         node["status"] = FileStatus.missing()
         finished_nodes = {node for node, attrs in self.import_graph.nodes(data=True)
-                          if attrs.get("status") == FileStatus.yes()}
+                          if attrs.get("status").ported}
         # tag nodes that have finished parents, depth of 1
         for node in finished_nodes:
             # does not get root nodes because they are not at end of an out_edge
@@ -1057,19 +1060,33 @@ def yaml_md_load(wikicontent: bytes):
                 parents = {parent for parent, _ in self.import_graph.in_edges(target)}
                 if parents.issubset(finished_nodes):
                     target_node = self.import_graph.nodes[target]
-                    if not target_node.get("status"):
-                        target_node["status"] = FileStatus.ready()
+                    target_node["status"].comments = "ready"
         # now to get root nodes
         for target, degree in self.import_graph.in_degree():
             target_node = self.import_graph.nodes[target]
-            if degree > 0 or target_node.get("status"):
+            if degree > 0 or target_node["status"].comments:
                 continue
-            target_node["status"] = FileStatus.ready()
-        for _, node in self.import_graph.nodes(data=True):
+            target_node["status"].comments = "ready"
+        for node_name, node in self.import_graph.nodes(data=True):
             if not node.get("status"):
                 continue
             node["style"] = "filled"
-            node["fillcolor"] = node["status"].color
+            node["fillcolor"] = self.status_color(node["status"])
+            if node_name in existing_files:
+                node["color"] = "red"
+
+    @staticmethod
+    def status_color(status: FileStatus) -> Optional[str]:
+        """
+        How each status should be colored, should be applied to the node attrs at "fillcolor".
+        """
+        if status.ported:
+            return "green"
+        if status.mathlib4_pr:
+            return "lightskyblue"
+        if status.comments and "ready" in status.comments:
+            return "turquoise1"
+        return None
 
     def modules_used(self, module: str) -> List[str]:
         """