Add pybind11-based glue code for new typegraph metrics.

PiperOrigin-RevId: 336195189

pytype/CMakeLists.txt

@@ -277,17 +277,15 @@ py_library(
       .utils
 )
 
-# This library is disabled because cfg.py can conflict with the typegraph.cfg
-# extension module.
-# py_library(
-#   NAME
-#     cfg_py
-#   SRCS
-#     typegraph/cfg.py
-#   DEPS
-#     .debug
-#     .metrics
-# )
+py_library(
+  NAME
+    cfg_py
+  SRCS
+    typegraph/cfg.py
+  DEPS
+    .debug
+    .metrics
+)
 
 py_library(
   NAME
@@ -660,6 +658,17 @@ py_test(
     .pytd
 )
 
+py_test(
+  NAME
+    typegraph_metrics_test
+  SRCS
+    typegraph_metrics_test.py
+  DEPS
+    .config
+    .libvm
+    pytype.typegraph.cfg
+    pytype.tests.test_base
+
 add_subdirectory(overlays)
 add_subdirectory(pyc)
 add_subdirectory(pyi)

pytype/typegraph/cfg.cc

@@ -1,3 +1,5 @@
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 #include <Python.h>
 #include <structseq.h>
 
@@ -6,12 +8,11 @@
 #include <utility>
 
 #include "cfg_logging.h"
+#include "metrics.h"
 #include "typegraph.h"
 
 namespace typegraph = devtools_python_typegraph;
 
-static const char* kModuleName = "cfg";
-
 #if PY_MAJOR_VERSION >= 3
 #  define PyString_Check(s) (PyBytes_Check(s) || PyUnicode_Check(s))
 #  define PyString_FromString PyUnicode_FromString
@@ -477,13 +478,22 @@ static PyObject* is_reachable(PyProgramObj* self,
   }
 }
 
+PyDoc_STRVAR(calculate_metrics_doc, "Get a snapshot of the program's metrics.");
+
+static PyObject* calculate_metrics(PyProgramObj* self, PyObject* _args) {
+  auto data = self->program->CalculateMetrics();
+  return pybind11::cast(data).release().ptr();
+}
+
 static PyMethodDef program_methods[] = {
   {"NewCFGNode", reinterpret_cast<PyCFunction>(NewCFGNode),
     METH_VARARGS|METH_KEYWORDS, new_cfg_node_doc},
   {"NewVariable", reinterpret_cast<PyCFunction>(NewVariable),
     METH_VARARGS|METH_KEYWORDS, new_variable_doc},
   {"is_reachable", reinterpret_cast<PyCFunction>(is_reachable),
    METH_VARARGS|METH_KEYWORDS, is_reachable_doc},
+  {"calculate_metrics", reinterpret_cast<PyCFunction>(calculate_metrics),
+   METH_NOARGS, calculate_metrics_doc},
   {0, 0, 0, nullptr}  // sentinel
 };
 
@@ -1438,34 +1448,7 @@ PyTypeObject PyVariable = {
   tp_methods : variable_methods,
 };
 
-// --- typegraph ---------------------------------------------------------------
-
-PyDoc_STRVAR(typegraph_doc,
-    "Typegraph is a points-to / dataflow / cfg graph engine.\n"
-    "It can be used to run reaching-definition queries on a nested CFG graph "
-    "and to model path-specific visibility of nested data structures.");
-
-static PyMethodDef typegraph_methods[] = {
-  {0, 0, 0, nullptr},  // sentinel
-};
-
-#if PY_MAJOR_VERSION >= 3
-static struct PyModuleDef typegraph_moduledef = {
-  PyModuleDef_HEAD_INIT,
-  m_name : kModuleName,
-  m_doc : typegraph_doc,
-  m_size : -1,
-  m_methods : typegraph_methods,
-#if PY_VERSION_HEX >= 0x03050000  // 3.5
-  m_slots : nullptr,
-#else
-  m_reload : nullptr,
-#endif
-  m_traverse : nullptr,
-  m_clear : &pytype::typegraph::internal::CFGLogger::Shutdown,
-  m_free : nullptr,
-};
-#endif
+// --- cfg module and metrics --------------------------------------------------
 
 static PyObject* InitModule(PyObject* module) {
   PyObject* module_dict = PyModule_GetDict(module);
@@ -1529,21 +1512,67 @@ static PyObject* InitModule(PyObject* module) {
   return module;
 }
 
-#if PY_MAJOR_VERSION >= 3
-PyMODINIT_FUNC PyInit_cfg(void) {
-  PyObject* module = PyModule_Create(&typegraph_moduledef);
-  pytype::typegraph::internal::CFGLogger::Init();
-  return InitModule(module);
-}
-#else
-PyMODINIT_FUNC initcfg(void) {
-  PyObject* module = Py_InitModule3(
-      kModuleName, typegraph_methods, typegraph_doc);
+// This creates a module called cfg inside typegraph.
+// The full path is pytype.typegraph.cfg.
+PYBIND11_MODULE(cfg, m) {
+  m.doc() = "Typegraph is a points-to / dataflow / cfg graph engine.\n"
+    "It can be used to run reaching-definition queries on a nested CFG graph "
+    "and to model path-specific visibility of nested data structures.";
+
+  pybind11::class_<typegraph::NodeMetrics>(m, "NodeMetrics")
+      .def_property_readonly("incoming_edge_count",
+                             &typegraph::NodeMetrics::incoming_edge_count)
+      .def_property_readonly("outgoing_edge_count",
+                             &typegraph::NodeMetrics::outgoing_edge_count)
+      .def_property_readonly("has_condition",
+                             &typegraph::NodeMetrics::has_condition);
+
+  pybind11::class_<typegraph::VariableMetrics>(m, "VariableMetrics")
+      .def_property_readonly("binding_count",
+                             &typegraph::VariableMetrics::binding_count)
+      .def_property_readonly("node_ids", &typegraph::VariableMetrics::node_ids);
+
+  pybind11::class_<typegraph::QueryMetrics>(m, "QueryMetrics")
+      .def_property_readonly("nodes_visited",
+                             &typegraph::QueryMetrics::nodes_visited)
+      .def_property_readonly("start_node", &typegraph::QueryMetrics::start_node)
+      .def_property_readonly("end_node", &typegraph::QueryMetrics::end_node)
+      .def_property_readonly("initial_binding_count",
+                             &typegraph::QueryMetrics::initial_binding_count)
+      .def_property_readonly("total_binding_count",
+                             &typegraph::QueryMetrics::total_binding_count)
+      .def_property_readonly("shortcircuited",
+                             &typegraph::QueryMetrics::shortcircuited)
+      .def_property_readonly("from_cache",
+                             &typegraph::QueryMetrics::from_cache);
+
+  pybind11::class_<typegraph::CacheMetrics>(m, "CacheMetrics")
+      .def_property_readonly("total_size", &typegraph::CacheMetrics::total_size)
+      .def_property_readonly("hits", &typegraph::CacheMetrics::hits)
+      .def_property_readonly("misses", &typegraph::CacheMetrics::misses);
+
+  pybind11::class_<typegraph::SolverMetrics>(m, "SolverMetrics")
+      .def_property_readonly("query_metrics",
+                             &typegraph::SolverMetrics::query_metrics)
+      .def_property_readonly("cache_metrics",
+                             &typegraph::SolverMetrics::cache_metrics);
+
+  pybind11::class_<typegraph::Metrics>(m, "Metrics")
+      .def_property_readonly("binding_count",
+                             &typegraph::Metrics::binding_count)
+      .def_property_readonly("cfg_node_metrics",
+                             &typegraph::Metrics::cfg_node_metrics)
+      .def_property_readonly("variable_metrics",
+                             &typegraph::Metrics::variable_metrics)
+      .def_property_readonly("solver_metrics",
+                             &typegraph::Metrics::solver_metrics)
+      .def_property_readonly("reachability_metrics",
+                             &typegraph::Metrics::reachability_metrics);
+
   PyType_Ready(&PyProgram);
   PyType_Ready(&PyCFGNode);
   PyType_Ready(&PyVariable);
   PyType_Ready(&PyBinding);
-  InitModule(module);
+  InitModule(m.ptr());
   pytype::typegraph::internal::CFGLogger::Init();
 }
-#endif

pytype/typegraph/typegraph_test.cc

@@ -237,5 +237,31 @@ TEST_F(TypeGraphTest, TestBindingIDs) {
   EXPECT_EQ(1, ax2->id());
   EXPECT_EQ(2, p.next_binding_id());
 }
+
+
+TEST_F(TypeGraphTest, TestMetrics) {
+  // Small test for making sure Metrics are collected.
+
+  Program p;
+  Variable* x = p.NewVariable();
+  CFGNode* n0 = p.NewCFGNode("n0");
+  CFGNode* n1 = n0->ConnectNew("n1");
+  int one = 1;
+  Binding* ax1 = AddBinding(x, &one, n1, {});
+  ax1->AddOrigin(n0);
+
+  auto metrics = p.CalculateMetrics();
+  EXPECT_EQ(metrics.binding_count(), 1);
+
+  auto cfgm = metrics.cfg_node_metrics();
+  EXPECT_EQ(cfgm[0].incoming_edge_count(), 0);
+  EXPECT_EQ(cfgm[0].outgoing_edge_count(), 1);
+  EXPECT_FALSE(cfgm[0].has_condition());
+
+  auto varm = metrics.variable_metrics();
+  EXPECT_EQ(varm.size(), 1);
+  EXPECT_EQ(varm[0].binding_count(), 1);
+  EXPECT_THAT(varm[0].node_ids(), testing::UnorderedElementsAre(0, 1));
+}
 }  // namespace
 }  // namespace devtools_python_typegraph

pytype/typegraph_metrics_test.py

@@ -0,0 +1,41 @@
+"""Basic tests for accessing typegraph metrics from Python."""
+import textwrap
+
+from pytype import analyze
+from pytype import errors
+from pytype import typegraph
+from pytype.tests import test_base
+
+
+class MetricsTest(test_base.BaseTest):
+
+  def setUp(self):
+    super().setUp()
+    self.errorlog = errors.ErrorLog()
+    self.vm = analyze.CallTracer(self.errorlog, self.options, self.loader)
+
+  def run_program(self, src):
+    return self.vm.run_program(textwrap.dedent(src), "", maximum_depth=10)
+
+  def assertNotEmpty(self, container, msg=None):
+    if not container:
+      self.fail("{!r} has length of 0.".format(container), msg)
+
+  def test_basics(self):
+    self.run_program("""
+        def foo(x: str) -> int:
+          return x + 1
+        a = foo(1)
+    """)
+    metrics = self.vm.program.calculate_metrics()
+    # No specific numbers are used to prevent this from being a change detector.
+    self.assertIsInstance(metrics, typegraph.cfg.Metrics)
+    self.assertGreater(metrics.binding_count, 0)
+    self.assertNotEmpty(metrics.cfg_node_metrics)
+    self.assertNotEmpty(metrics.variable_metrics)
+    # TODO(tsudol): Need to implement solver and reachability metrics.
+    # self.assertNotEmpty(metrics.solver_metrics)
+    self.assertGreater(metrics.reachability_metrics.total_size, 0)
+
+
+test_base.main(globals(), __name__ == "__main__")