# Copyright 2022 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Script to generate header cc and unittest file for a class in chromium."""
_DOCUMENTATION = r"""Usage:
To generate default model template files:
python3 components/segmentation_platform/internal/tools/create_class.py \
--segment_id MY_FEATURE_USER
To generate generic header and cc files:
python3 components/segmentation_platform/internal/tools/create_class.py \
--header src/dir/class_name.h
If any of the file already exists then prints a log and does not touch the
file, but still creates the remaining files.
"""
import argparse
import datetime
import logging
import os
import sys
_HEADER_TEMPLATE = """// Copyright {year} The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#ifndef {macro}
#define {macro}
namespace {namespace} {{
class {clas} {{
public:
{clas}();
~{clas}();
{clas}(const {clas}&) = delete;
{clas}& operator=(const {clas}&) = delete;
private:
}};
}}
#endif // {macro}
"""
_CC_TEMPLATE = """// Copyright {year} The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#include "{file_path}"
namespace {namespace} {{
{clas}::{clas} () = default;
{clas}::~{clas}() = default;
}}
"""
_TEST_TEMPLATE = """// Copyright {year} The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#include "{file_path}"
#include "testing/gtest/include/gtest/gtest.h"
namespace {namespace} {{
class {test_class} : public testing::Test {{
public:
{test_class}() = default;
~{test_class}() override = default;
void SetUp() override {{
Test::SetUp();
}}
void TearDown() override {{
Test::TearDown();
}}
protected:
}};
TEST_F({test_class}, Test) {{
}}
}}
"""
_MODEL_HEADER_TEMPLATE = """// Copyright {year} The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#ifndef {macro}
#define {macro}
#include <memory>
#include "base/feature_list.h"
#include "components/segmentation_platform/public/config.h"
#include "components/segmentation_platform/public/model_provider.h"
namespace {namespace} {{
// Feature flag for enabling {clas} segment.
BASE_DECLARE_FEATURE(kSegmentationPlatform{clas});
// Model to predict whether the user belongs to {clas} segment.
class {clas} : public DefaultModelProvider {{
public:
static constexpr char k{clas}Key[] = "{segmentation_key}";
static constexpr char k{clas}UmaName[] = "{clas}";
{clas}();
~{clas}() override = default;
{clas}(const {clas}&) = delete;
{clas}& operator=(const {clas}&) = delete;
static std::unique_ptr<Config> GetConfig();
// ModelProvider implementation.
std::unique_ptr<ModelConfig> GetModelConfig() override;
void ExecuteModelWithInput(const ModelProvider::Request& inputs,
ExecutionCallback callback) override;
}};
}}
#endif // {macro}
"""
_MODEL_CC_TEMPLATE = """// Copyright {year} The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#include "{file_path}"
#include <memory>
#include "base/task/sequenced_task_runner.h"
#include "components/segmentation_platform/internal/metadata/metadata_writer.h"
#include "components/segmentation_platform/public/config.h"
#include "components/segmentation_platform/public/proto/aggregation.pb.h"
#include "components/segmentation_platform/public/proto/model_metadata.pb.h"
namespace {namespace} {{
BASE_FEATURE(kSegmentationPlatform{clas},
"SegmentationPlatform{clas}",
base::FEATURE_DISABLED_BY_DEFAULT);
namespace {{
using proto::SegmentId;
// Default parameters for {clas} model.
constexpr SegmentId kSegmentId = SegmentId::{segment_id};
constexpr int64_t kModelVersion = 1;
// Store 28 buckets of input data (28 days).
constexpr int64_t kSignalStorageLength = 28;
// Wait until we have 7 days of data.
constexpr int64_t kMinSignalCollectionLength = 7;
// Refresh the result every 7 days.
constexpr int64_t kResultTTLDays = 7;
// InputFeatures.
// Enum values for the Example.EnumHistogram.
constexpr std::array<int32_t, 3> kEnumValues{{
0, 3, 4
}};
// Set UMA metrics to use as input.
// TODO: Fill in the necessary signals for prediction.
constexpr std::array<MetadataWriter::UMAFeature, 3> kUMAFeatures = {{
// Total amount of times user action was recorded in last 14 days.
MetadataWriter::UMAFeature::FromUserAction("UserActionName", 14),
// Total value of all records of the histogram in last 7 days.
MetadataWriter::UMAFeature::FromValueHistogram(
"Example.ValueHistogram", 7, proto::Aggregation::SUM),
// Total count of number of records of enum histogram with given values.
MetadataWriter::UMAFeature::FromEnumHistogram(
"Example.EnumHistogram",
14,
kEnumValues.data(),
kEnumValues.size()),
}};
}} // namespace
// static
std::unique_ptr<Config> {clas}::GetConfig() {{
if (!base::FeatureList::IsEnabled(
kSegmentationPlatform{clas})) {{
return nullptr;
}}
auto config = std::make_unique<Config>();
config->segmentation_key = k{clas}Key;
config->segmentation_uma_name = k{clas}UmaName;
config->AddSegmentId(kSegmentId,
std::make_unique<{clas}>());
config->auto_execute_and_cache = false;
return config;
}}
{clas}::{clas}()
: DefaultModelProvider(kSegmentId) {{}}
std::unique_ptr<DefaultModelProvider::ModelConfig> {clas}::GetModelConfig() {{
proto::SegmentationModelMetadata metadata;
MetadataWriter writer(&metadata);
writer.SetDefaultSegmentationMetadataConfig(
kMinSignalCollectionLength,
kSignalStorageLength);
// Set output config.
const char kNot{clas}Label[] = "Not{clas}";
writer.AddOutputConfigForBinaryClassifier(
0.5,
/*positive_label=*/k{clas}UmaName,
kNot{clas}Label);
writer.AddPredictedResultTTLInOutputConfig(
/*top_label_to_ttl_list=*/{{}},
/*default_ttl=*/kResultTTLDays, proto::TimeUnit::DAY);
// Set features.
writer.AddUmaFeatures(kUMAFeatures.data(),
kUMAFeatures.size());
return std::make_unique<ModelConfig>(std::move(metadata), kModelVersion);
}}
void {clas}::ExecuteModelWithInput(
const ModelProvider::Request& inputs,
ExecutionCallback callback) {{
// Invalid inputs.
if (inputs.size() != kUMAFeatures.size()) {{
base::SequencedTaskRunner::GetCurrentDefault()->PostTask(
FROM_HERE, base::BindOnce(std::move(callback), std::nullopt));
return;
}}
// TODO: Update the heuristics here to return 1 when the user belongs to
// {clas}.
float result = 0;
const int user_action_count = inputs[0];
const int value_histogram_total = inputs[1];
const int enum_hit_count = inputs[2];
if (user_action_count && value_histogram_total && enum_hit_count)
result = 1;
base::SequencedTaskRunner::GetCurrentDefault()->PostTask(
FROM_HERE,
base::BindOnce(std::move(callback), ModelProvider::Response(1, result)));
}}
}}
"""
_MODEL_TEST_TEMPLATE = """// Copyright {year} The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#include "{file_path}"
#include "components/segmentation_platform/embedder/default_model/default_model_test_base.h"
#include "testing/gtest/include/gtest/gtest.h"
namespace {namespace} {{
class {test_class} : public DefaultModelTestBase {{
public:
{test_class}() : DefaultModelTestBase(std::make_unique<{clas}>()) {{}}
~{test_class}() override = default;
}};
TEST_F({test_class}, InitAndFetchModel) {{
ExpectInitAndFetchModel();
}}
TEST_F({test_class}, ExecuteModelWithInput) {{
// TODO: Add test cases to verify if the heuristic returns the right segment.
ExpectExecutionWithInput(/*inputs=*/{{1, 2, 3}}, /*expected_error=*/false,
/*expected_result=*/{{1}});
}}
}}
"""
def _GetLogger():
"""Logger for the tool."""
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger('create_class')
logger.setLevel(level=logging.INFO)
return logger
def _WriteFile(path, type_str, contents):
"""Writes a file with contents to the path, if not exists."""
if os.path.exists(path):
_GetLogger().error('%s already exists', type_str)
return
_GetLogger().info('Writing %s file %s', type_str, path)
with open(path, 'w') as f:
f.write(contents)
def _GetClassNameFromFile(header):
"""Gets a class name from the header file name."""
file_base = os.path.basename(header).replace('.h', '')
class_name = ''
for i in range(len(file_base)):
if i == 0 or file_base[i - 1] == '_':
class_name += file_base[i].upper()
elif file_base[i] == '_':
continue
else:
class_name += file_base[i]
return class_name
def _GetSegmentationKeyFromFile(header):
"""Gets the segmentation key based on the header file."""
return os.path.basename(header).replace('.h', '')
def _GetHeader(args):
"""Parses the args and returns path to the header file."""
if args.header:
if '.h' not in args.header:
raise ValueError('The first argument should be a path to header')
_GetLogger().info('Creating class for header %s', args.header)
return args.header
if args.segment_id:
_PREFIXES_TO_REMOVE = [
'OPTIMIZATION_TARGET_SEGMENTATION_', 'OPTIMIZATION_TARGET_'
]
_GetLogger().info('Creating default model for %s', args.segment_id)
model_name = args.segment_id
for prefix in _PREFIXES_TO_REMOVE:
print(prefix, model_name, model_name.startswith(prefix))
if model_name.startswith(prefix):
model_name = model_name[len(prefix):]
break
print(model_name)
return (
'components/segmentation_platform/embedder/default_model/%s.h' %
model_name.lower())
raise ValueError('Required either --header or --segment_id argument.')
def _CreateFilesForClass(args):
"""Creates header cc and test files for the class."""
header_template = _HEADER_TEMPLATE
cc_template = _CC_TEMPLATE
test_template = _TEST_TEMPLATE
if args.segment_id:
header_template = _MODEL_HEADER_TEMPLATE
cc_template = _MODEL_CC_TEMPLATE
test_template = _MODEL_TEST_TEMPLATE
header = _GetHeader(args)
file_cc = header.replace('.h', '.cc')
file_test = header.replace('.h', '_unittest.cc')
format_args = {}
format_args['year'] = datetime.date.today().year
format_args['file_path'] = header
format_args['macro'] = (
header.replace('/', '_').replace('.', '_').upper() + '_')
format_args['clas'] = _GetClassNameFromFile(header)
format_args['segment_id'] = args.segment_id
format_args['segmentation_key'] = _GetSegmentationKeyFromFile(header)
format_args['namespace'] = args.namespace
format_args['test_class'] = format_args['clas'] + 'Test'
contents = header_template.format_map(format_args)
_WriteFile(header, 'Header', contents)
contents = cc_template.format_map(format_args)
_WriteFile(file_cc, 'CC', contents)
contents = test_template.format_map(format_args)
_WriteFile(file_test, 'Test', contents)
def _CreateOptionParser():
"""Options parser for the tool."""
parser = argparse.ArgumentParser(
description=_DOCUMENTATION,
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument('--header',
help='Path to the header file from src/',
default='')
parser.add_argument('--segment_id',
help='The segment ID enum value',
default='')
parser.add_argument('--namespace',
dest='namespace',
default='segmentation_platform')
return parser
def main():
parser = _CreateOptionParser()
args = parser.parse_args()
_CreateFilesForClass(args)
if __name__ == '__main__':
sys.exit(main())
Codebase Browser
chromium